Address linter issues

Signed-off-by: Maysun J Faisal <[email protected]>

Author: maysunfaisal

src/app/endpoints/conversations_v2.py

@@ -314,19 +314,13 @@ def check_conversation_existence(user_id: str, conversation_id: str) -> None:
 
 def transform_chat_message(entry: CacheEntry) -> dict[str, Any]:
     """Transform the message read from cache into format used by response payload."""
-    user_message = {
-        "content": entry.query, 
-        "type": "user"
-    }
-    assistant_message: dict[str, Any] = {
-        "content": entry.response,
-        "type": "assistant"
-    }
+    user_message = {"content": entry.query, "type": "user"}
+    assistant_message: dict[str, Any] = {"content": entry.response, "type": "assistant"}
 
     # If referenced_documents exist on the entry, add them to the assistant message
     if entry.referenced_documents is not None:
         assistant_message["referenced_documents"] = [
-            doc.model_dump(mode='json') for doc in entry.referenced_documents
+            doc.model_dump(mode="json") for doc in entry.referenced_documents
         ]
 
     return {

src/app/endpoints/query.py

@@ -22,7 +22,6 @@
 from llama_stack_client.types.model_list_response import ModelListResponse
 from llama_stack_client.types.shared.interleaved_content_item import TextContentItem
 from llama_stack_client.types.tool_execution_step import ToolExecutionStep
-from pydantic import AnyUrl
 
 import constants
 import metrics
@@ -341,9 +340,9 @@ async def query_endpoint_handler(  # pylint: disable=R0914
             model=model_id,
             started_at=started_at,
             completed_at=completed_at,
-            referenced_documents=referenced_documents if referenced_documents else None
+            referenced_documents=referenced_documents if referenced_documents else None,
         )
-            
+
         store_conversation_into_cache(
             configuration,
             user_id,

src/app/endpoints/streaming_query.py

@@ -21,7 +21,6 @@
 )
 from llama_stack_client.types.shared import ToolCall
 from llama_stack_client.types.shared.interleaved_content_item import TextContentItem
-from pydantic import AnyUrl
 
 from app.database import get_session
 from app.endpoints.query import (
@@ -48,7 +47,7 @@
 from models.config import Action
 from models.database.conversations import UserConversation
 from models.requests import QueryRequest
-from models.responses import ForbiddenResponse, UnauthorizedResponse, ReferencedDocument
+from models.responses import ForbiddenResponse, UnauthorizedResponse
 from utils.endpoints import (
     check_configuration_loaded,
     create_referenced_documents_with_metadata,
@@ -867,7 +866,9 @@ async def response_generator(
 
             completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
 
-            referenced_documents = create_referenced_documents_with_metadata(summary, metadata_map)
+            referenced_documents = create_referenced_documents_with_metadata(
+                summary, metadata_map
+            )
 
             cache_entry = CacheEntry(
                 query=query_request.query,
@@ -876,9 +877,11 @@ async def response_generator(
                 model=model_id,
                 started_at=started_at,
                 completed_at=completed_at,
-                referenced_documents=referenced_documents if referenced_documents else None
+                referenced_documents=(
+                    referenced_documents if referenced_documents else None
+                ),
             )
-            
+
             store_conversation_into_cache(
                 configuration,
                 user_id,

src/cache/postgres_cache.py

@@ -219,7 +219,9 @@ def get(
                 docs_data = conversation_entry[6]
                 docs_obj = None
                 if docs_data:
-                    docs_obj = [ReferencedDocument.model_validate(doc) for doc in docs_data]
+                    docs_obj = [
+                        ReferencedDocument.model_validate(doc) for doc in docs_data
+                    ]
                 cache_entry = CacheEntry(
                     query=conversation_entry[0],
                     response=conversation_entry[1],
@@ -257,7 +259,10 @@ def insert_or_append(
         try:
             referenced_documents_json = None
             if cache_entry.referenced_documents:
-                docs_as_dicts = [doc.model_dump(mode='json') for doc in cache_entry.referenced_documents]
+                docs_as_dicts = [
+                    doc.model_dump(mode="json")
+                    for doc in cache_entry.referenced_documents
+                ]
                 referenced_documents_json = json.dumps(docs_as_dicts)
 
             # the whole operation is run in one transaction

src/cache/sqlite_cache.py

@@ -257,7 +257,9 @@ def insert_or_append(
 
         referenced_documents_json = None
         if cache_entry.referenced_documents:
-            docs_as_dicts = [doc.model_dump(mode='json') for doc in cache_entry.referenced_documents]
+            docs_as_dicts = [
+                doc.model_dump(mode="json") for doc in cache_entry.referenced_documents
+            ]
             referenced_documents_json = json.dumps(docs_as_dicts)
 
         cursor.execute(

src/models/cache_entry.py

@@ -1,7 +1,7 @@
 """Model for conversation history cache entry."""
 
-from pydantic import BaseModel, Field
 from typing import List
+from pydantic import BaseModel
 from models.responses import ReferencedDocument
 
 

src/models/responses.py

@@ -161,6 +161,7 @@ class ToolCall(BaseModel):
     arguments: dict[str, Any] = Field(description="Arguments passed to the tool")
     result: Optional[dict[str, Any]] = Field(None, description="Result from the tool")
 
+
 class ConversationData(BaseModel):
     """Model representing conversation data returned by cache list operations.
 
@@ -174,6 +175,7 @@ class ConversationData(BaseModel):
     topic_summary: str | None
     last_message_timestamp: float
 
+
 class ReferencedDocument(BaseModel):
     """Model representing a document referenced in generating a response.
 
@@ -186,9 +188,7 @@ class ReferencedDocument(BaseModel):
         None, description="URL of the referenced document"
     )
 
-    doc_title: str | None = Field(
-        None, description="Title of the referenced document"
-    )
+    doc_title: str | None = Field(None, description="Title of the referenced document")
 
 
 class QueryResponse(BaseModel):

tests/unit/app/endpoints/test_conversations_v2.py

@@ -75,31 +75,33 @@ def test_transform_message_without_documents(self) -> None:
             # referenced_documents is None by default
         )
         transformed = transform_chat_message(entry)
-        
+
         assistant_message = transformed["messages"][1]
 
         # Assert that the key is NOT present when the list is None
         assert "referenced_documents" not in assistant_message
 
     def test_transform_message_with_referenced_documents(self) -> None:
         """Test the transformation when referenced_documents are present."""
-        docs = [ReferencedDocument(doc_title="Test Doc", doc_url=AnyUrl("http://example.com"))]
-        
+        docs = [
+            ReferencedDocument(
+                doc_title="Test Doc", doc_url=AnyUrl("http://example.com")
+            )
+        ]
         entry = CacheEntry(
             query="query",
             response="response",
             provider="provider",
             model="model",
             started_at="2024-01-01T00:00:00Z",
             completed_at="2024-01-01T00:00:05Z",
-            referenced_documents=docs
+            referenced_documents=docs,
         )
 
         transformed = transform_chat_message(entry)
         assistant_message = transformed["messages"][1]
-        
+
         assert "referenced_documents" in assistant_message
-        
         ref_docs = assistant_message["referenced_documents"]
         assert len(ref_docs) == 1
         assert ref_docs[0]["doc_title"] == "Test Doc"
@@ -114,7 +116,7 @@ def test_transform_message_with_empty_referenced_documents(self) -> None:
             model="model",
             started_at="2024-01-01T00:00:00Z",
             completed_at="2024-01-01T00:00:05Z",
-            referenced_documents=[] # Explicitly empty
+            referenced_documents=[],  # Explicitly empty
         )
 
         transformed = transform_chat_message(entry)

tests/unit/app/endpoints/test_query.py

@@ -167,6 +167,7 @@ def test_is_transcripts_disabled(setup_configuration, mocker) -> None:
     assert is_transcripts_enabled() is False, "Transcripts should be disabled"
 
 
+# pylint: disable=too-many-locals
 async def _test_query_endpoint_handler(
     mocker, dummy_request: Request, store_transcript_to_file=False
 ) -> None:
@@ -184,12 +185,16 @@ async def _test_query_endpoint_handler(
         store_transcript_to_file
     )
     mocker.patch("app.endpoints.query.configuration", mock_config)
-    
-    mock_store_in_cache = mocker.patch("app.endpoints.query.store_conversation_into_cache")
+
+    mock_store_in_cache = mocker.patch(
+        "app.endpoints.query.store_conversation_into_cache"
+    )
 
     # Create mock referenced documents to simulate a successful RAG response
     mock_referenced_documents = [
-        ReferencedDocument(doc_title="Test Doc 1", doc_url=AnyUrl("http://example.com/1"))
+        ReferencedDocument(
+            doc_title="Test Doc 1", doc_url=AnyUrl("http://example.com/1")
+        )
     ]
 
     summary = TurnSummary(
@@ -208,7 +213,12 @@ async def _test_query_endpoint_handler(
 
     mocker.patch(
         "app.endpoints.query.retrieve_response",
-        return_value=(summary, conversation_id, mock_referenced_documents, TokenCounter()),
+        return_value=(
+            summary,
+            conversation_id,
+            mock_referenced_documents,
+            TokenCounter(),
+        ),
     )
     mocker.patch(
         "app.endpoints.query.select_model_and_provider_id",
@@ -237,11 +247,12 @@ async def _test_query_endpoint_handler(
     # Assert the response is as expected
     assert response.response == summary.llm_response
     assert response.conversation_id == conversation_id
-    
+
     # Assert that mock was called and get the arguments
     mock_store_in_cache.assert_called_once()
     call_args = mock_store_in_cache.call_args[0]
-    # Extract CacheEntry object from the call arguments, it's the 4th argument from the func signature
+    # Extract CacheEntry object from the call arguments,
+    # it's the 4th argument from the func signature
     cached_entry = call_args[3]
 
     assert isinstance(cached_entry, CacheEntry)

tests/unit/app/endpoints/test_streaming_query.py

@@ -209,6 +209,7 @@ async def test_streaming_query_endpoint_on_connection_error(mocker):
     assert response.media_type == "text/event-stream"
 
 
+# pylint: disable=too-many-locals
 async def _test_streaming_query_endpoint_handler(mocker, store_transcript=False):
     """Test the streaming query endpoint handler."""
     mock_client = mocker.AsyncMock()
@@ -297,8 +298,9 @@ async def _test_streaming_query_endpoint_handler(mocker, store_transcript=False)
         ),
     ]
 
-    mock_store_in_cache = mocker.patch("app.endpoints.streaming_query.store_conversation_into_cache")
-
+    mock_store_in_cache = mocker.patch(
+        "app.endpoints.streaming_query.store_conversation_into_cache"
+    )
     query = "What is OpenStack?"
     mocker.patch(
         "app.endpoints.streaming_query.retrieve_response",
@@ -359,20 +361,23 @@ async def _test_streaming_query_endpoint_handler(mocker, store_transcript=False)
     referenced_documents = d["data"]["referenced_documents"]
     assert len(referenced_documents) == 2
     assert referenced_documents[1]["doc_title"] == "Doc2"
-    
+
     # Assert that mock was called and get the arguments
     mock_store_in_cache.assert_called_once()
     call_args = mock_store_in_cache.call_args[0]
-    # Extract CacheEntry object from the call arguments, it's the 4th argument from the func signature
+    # Extract CacheEntry object from the call arguments,
+    # it's the 4th argument from the func signature
     cached_entry = call_args[3]
-    
+
     # Assert that the CacheEntry was constructed correctly
     assert isinstance(cached_entry, CacheEntry)
     assert cached_entry.response == "LLM answer"
     assert cached_entry.referenced_documents is not None
     assert len(cached_entry.referenced_documents) == 2
     assert cached_entry.referenced_documents[0].doc_title == "Doc1"
-    assert str(cached_entry.referenced_documents[1].doc_url) == "https://example.com/doc2"
+    assert (
+        str(cached_entry.referenced_documents[1].doc_url) == "https://example.com/doc2"
+    )
 
     # Assert the store_transcript function is called if transcripts are enabled
     if store_transcript: