Skip to content

Commit c871aaa

Browse files
maysunfaisalmaysunfaisal
committed
Address linter issues
Signed-off-by: Maysun J Faisal <[email protected]>
1 parent ef0b934 commit c871aaa

File tree

13 files changed

+234
-120
lines changed

13 files changed

+234
-120
lines changed

‎docs/openapi.json‎

Lines changed: 100 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1787,6 +1787,9 @@
17871787
},
17881788
"type": "array",
17891789
"title": "Byok Rag"
1790+
},
1791+
"quota_handlers": {
1792+
"$ref": "#/components/schemas/QuotaHandlersConfiguration"
17901793
}
17911794
},
17921795
"additionalProperties": false,
@@ -3590,6 +3593,103 @@
35903593
}
35913594
]
35923595
},
3596+
"QuotaHandlersConfiguration": {
3597+
"properties": {
3598+
"sqlite": {
3599+
"anyOf": [
3600+
{
3601+
"$ref": "#/components/schemas/SQLiteDatabaseConfiguration"
3602+
},
3603+
{
3604+
"type": "null"
3605+
}
3606+
]
3607+
},
3608+
"postgres": {
3609+
"anyOf": [
3610+
{
3611+
"$ref": "#/components/schemas/PostgreSQLDatabaseConfiguration"
3612+
},
3613+
{
3614+
"type": "null"
3615+
}
3616+
]
3617+
},
3618+
"limiters": {
3619+
"items": {
3620+
"$ref": "#/components/schemas/QuotaLimiterConfiguration"
3621+
},
3622+
"type": "array",
3623+
"title": "Limiters"
3624+
},
3625+
"scheduler": {
3626+
"$ref": "#/components/schemas/QuotaSchedulerConfiguration"
3627+
},
3628+
"enable_token_history": {
3629+
"type": "boolean",
3630+
"title": "Enable Token History",
3631+
"default": false
3632+
}
3633+
},
3634+
"additionalProperties": false,
3635+
"type": "object",
3636+
"title": "QuotaHandlersConfiguration",
3637+
"description": "Quota limiter configuration."
3638+
},
3639+
"QuotaLimiterConfiguration": {
3640+
"properties": {
3641+
"type": {
3642+
"type": "string",
3643+
"enum": [
3644+
"user_limiter",
3645+
"cluster_limiter"
3646+
],
3647+
"title": "Type"
3648+
},
3649+
"name": {
3650+
"type": "string",
3651+
"title": "Name"
3652+
},
3653+
"initial_quota": {
3654+
"type": "integer",
3655+
"minimum": 0.0,
3656+
"title": "Initial Quota"
3657+
},
3658+
"quota_increase": {
3659+
"type": "integer",
3660+
"minimum": 0.0,
3661+
"title": "Quota Increase"
3662+
},
3663+
"period": {
3664+
"type": "string",
3665+
"title": "Period"
3666+
}
3667+
},
3668+
"additionalProperties": false,
3669+
"type": "object",
3670+
"required": [
3671+
"type",
3672+
"name",
3673+
"initial_quota",
3674+
"quota_increase",
3675+
"period"
3676+
],
3677+
"title": "QuotaLimiterConfiguration",
3678+
"description": "Configuration for one quota limiter."
3679+
},
3680+
"QuotaSchedulerConfiguration": {
3681+
"properties": {
3682+
"period": {
3683+
"type": "integer",
3684+
"exclusiveMinimum": 0.0,
3685+
"title": "Period",
3686+
"default": 1
3687+
}
3688+
},
3689+
"type": "object",
3690+
"title": "QuotaSchedulerConfiguration",
3691+
"description": "Quota scheduler configuration."
3692+
},
35933693
"RAGChunk": {
35943694
"properties": {
35953695
"content": {
@@ -3973,45 +4073,6 @@
39734073
"title": "ToolsResponse",
39744074
"description": "Model representing a response to tools request."
39754075
},
3976-
"UnauthorizedResponse": {
3977-
"properties": {
3978-
"tools": {
3979-
"items": {
3980-
"additionalProperties": true,
3981-
"type": "object"
3982-
},
3983-
"type": "array",
3984-
"title": "Tools",
3985-
"description": "List of tools available from all configured MCP servers and built-in toolgroups",
3986-
"examples": [
3987-
[
3988-
{
3989-
"description": "Read contents of a file from the filesystem",
3990-
"identifier": "filesystem_read",
3991-
"parameters": [
3992-
{
3993-
"description": "Path to the file to read",
3994-
"name": "path",
3995-
"parameter_type": "string",
3996-
"required": true
3997-
}
3998-
],
3999-
"provider_id": "model-context-protocol",
4000-
"server_source": "http://localhost:3000",
4001-
"toolgroup_id": "filesystem-tools",
4002-
"type": "tool"
4003-
}
4004-
]
4005-
]
4006-
}
4007-
},
4008-
"type": "object",
4009-
"required": [
4010-
"tools"
4011-
],
4012-
"title": "ToolsResponse",
4013-
"description": "Model representing a response to tools request."
4014-
},
40154076
"UnauthorizedResponse": {
40164077
"properties": {
40174078
"detail": {

‎src/app/endpoints/conversations_v2.py‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -314,19 +314,13 @@ def check_conversation_existence(user_id: str, conversation_id: str) -> None:
314314

315315
def transform_chat_message(entry: CacheEntry) -> dict[str, Any]:
316316
"""Transform the message read from cache into format used by response payload."""
317-
user_message = {
318-
"content": entry.query,
319-
"type": "user"
320-
}
321-
assistant_message: dict[str, Any] = {
322-
"content": entry.response,
323-
"type": "assistant"
324-
}
317+
user_message = {"content": entry.query, "type": "user"}
318+
assistant_message: dict[str, Any] = {"content": entry.response, "type": "assistant"}
325319

326320
# If referenced_documents exist on the entry, add them to the assistant message
327321
if entry.referenced_documents is not None:
328322
assistant_message["referenced_documents"] = [
329-
doc.model_dump(mode='json') for doc in entry.referenced_documents
323+
doc.model_dump(mode="json") for doc in entry.referenced_documents
330324
]
331325

332326
return {

‎src/app/endpoints/query.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
from llama_stack_client.types.model_list_response import ModelListResponse
2323
from llama_stack_client.types.shared.interleaved_content_item import TextContentItem
2424
from llama_stack_client.types.tool_execution_step import ToolExecutionStep
25-
from pydantic import AnyUrl
2625

2726
import constants
2827
import metrics
@@ -341,9 +340,9 @@ async def query_endpoint_handler( # pylint: disable=R0914
341340
model=model_id,
342341
started_at=started_at,
343342
completed_at=completed_at,
344-
referenced_documents=referenced_documents if referenced_documents else None
343+
referenced_documents=referenced_documents if referenced_documents else None,
345344
)
346-
345+
347346
store_conversation_into_cache(
348347
configuration,
349348
user_id,

‎src/app/endpoints/streaming_query.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
)
2222
from llama_stack_client.types.shared import ToolCall
2323
from llama_stack_client.types.shared.interleaved_content_item import TextContentItem
24-
from pydantic import AnyUrl
2524

2625
from app.database import get_session
2726
from app.endpoints.query import (
@@ -48,7 +47,7 @@
4847
from models.config import Action
4948
from models.database.conversations import UserConversation
5049
from models.requests import QueryRequest
51-
from models.responses import ForbiddenResponse, UnauthorizedResponse, ReferencedDocument
50+
from models.responses import ForbiddenResponse, UnauthorizedResponse
5251
from utils.endpoints import (
5352
check_configuration_loaded,
5453
create_referenced_documents_with_metadata,
@@ -867,7 +866,9 @@ async def response_generator(
867866

868867
completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
869868

870-
referenced_documents = create_referenced_documents_with_metadata(summary, metadata_map)
869+
referenced_documents = create_referenced_documents_with_metadata(
870+
summary, metadata_map
871+
)
871872

872873
cache_entry = CacheEntry(
873874
query=query_request.query,
@@ -876,9 +877,11 @@ async def response_generator(
876877
model=model_id,
877878
started_at=started_at,
878879
completed_at=completed_at,
879-
referenced_documents=referenced_documents if referenced_documents else None
880+
referenced_documents=(
881+
referenced_documents if referenced_documents else None
882+
),
880883
)
881-
884+
882885
store_conversation_into_cache(
883886
configuration,
884887
user_id,

‎src/cache/postgres_cache.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -215,11 +215,13 @@ def get(
215215

216216
result = []
217217
for conversation_entry in conversation_entries:
218-
# Parse it back into an LLMResponse object
218+
# Parse referenced_documents back into ReferencedDocument objects
219219
docs_data = conversation_entry[6]
220220
docs_obj = None
221221
if docs_data:
222-
docs_obj = [ReferencedDocument.model_validate(doc) for doc in docs_data]
222+
docs_obj = [
223+
ReferencedDocument.model_validate(doc) for doc in docs_data
224+
]
223225
cache_entry = CacheEntry(
224226
query=conversation_entry[0],
225227
response=conversation_entry[1],
@@ -257,7 +259,10 @@ def insert_or_append(
257259
try:
258260
referenced_documents_json = None
259261
if cache_entry.referenced_documents:
260-
docs_as_dicts = [doc.model_dump(mode='json') for doc in cache_entry.referenced_documents]
262+
docs_as_dicts = [
263+
doc.model_dump(mode="json")
264+
for doc in cache_entry.referenced_documents
265+
]
261266
referenced_documents_json = json.dumps(docs_as_dicts)
262267

263268
# the whole operation is run in one transaction

‎src/cache/sqlite_cache.py‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -216,8 +216,18 @@ def get(
216216
docs_json_str = conversation_entry[6]
217217
docs_obj = None
218218
if docs_json_str:
219-
docs_data = json.loads(docs_json_str)
220-
docs_obj = [ReferencedDocument.model_validate(doc) for doc in docs_data]
219+
try:
220+
docs_data = json.loads(docs_json_str)
221+
docs_obj = [
222+
ReferencedDocument.model_validate(doc) for doc in docs_data
223+
]
224+
except (json.JSONDecodeError, ValueError) as e:
225+
logger.warning(
226+
"Failed to deserialize referenced_documents for "
227+
"conversation %s: %s",
228+
conversation_id,
229+
e,
230+
)
221231
cache_entry = CacheEntry(
222232
query=conversation_entry[0],
223233
response=conversation_entry[1],
@@ -257,8 +267,19 @@ def insert_or_append(
257267

258268
referenced_documents_json = None
259269
if cache_entry.referenced_documents:
260-
docs_as_dicts = [doc.model_dump(mode='json') for doc in cache_entry.referenced_documents]
261-
referenced_documents_json = json.dumps(docs_as_dicts)
270+
try:
271+
docs_as_dicts = [
272+
doc.model_dump(mode="json")
273+
for doc in cache_entry.referenced_documents
274+
]
275+
referenced_documents_json = json.dumps(docs_as_dicts)
276+
except (TypeError, ValueError) as e:
277+
logger.warning(
278+
"Failed to serialize referenced_documents for "
279+
"conversation %s: %s",
280+
conversation_id,
281+
e,
282+
)
262283

263284
cursor.execute(
264285
self.INSERT_CONVERSATION_HISTORY_STATEMENT,

‎src/models/cache_entry.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
"""Model for conversation history cache entry."""
22

3-
from pydantic import BaseModel, Field
4-
from typing import List
3+
from pydantic import BaseModel
54
from models.responses import ReferencedDocument
65

76

@@ -13,7 +12,7 @@ class CacheEntry(BaseModel):
1312
response: The response string
1413
provider: Provider identification
1514
model: Model identification
16-
additional_kwargs: additional property to store data like referenced documents
15+
referenced_documents: List of documents referenced by the response
1716
"""
1817

1918
query: str
@@ -22,4 +21,4 @@ class CacheEntry(BaseModel):
2221
model: str
2322
started_at: str
2423
completed_at: str
25-
referenced_documents: List[ReferencedDocument] | None = None
24+
referenced_documents: list[ReferencedDocument] | None = None

‎src/models/responses.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,7 @@ class ToolCall(BaseModel):
161161
arguments: dict[str, Any] = Field(description="Arguments passed to the tool")
162162
result: Optional[dict[str, Any]] = Field(None, description="Result from the tool")
163163

164+
164165
class ConversationData(BaseModel):
165166
"""Model representing conversation data returned by cache list operations.
166167
@@ -174,6 +175,7 @@ class ConversationData(BaseModel):
174175
topic_summary: str | None
175176
last_message_timestamp: float
176177

178+
177179
class ReferencedDocument(BaseModel):
178180
"""Model representing a document referenced in generating a response.
179181
@@ -186,9 +188,7 @@ class ReferencedDocument(BaseModel):
186188
None, description="URL of the referenced document"
187189
)
188190

189-
doc_title: str | None = Field(
190-
None, description="Title of the referenced document"
191-
)
191+
doc_title: str | None = Field(None, description="Title of the referenced document")
192192

193193

194194
class QueryResponse(BaseModel):

0 commit comments

Comments
 (0)