[AAQ-797] Refactor search base (#410)

* refactor search_base func

* cleanup json schema displays

Author: suzinyou

.secrets.baseline

@@ -231,15 +231,6 @@
         "line_number": 63
       }
     ],
-    "core_backend/app/question_answer/schemas.py": [
-      {
-        "type": "Secret Keyword",
-        "filename": "core_backend/app/question_answer/schemas.py",
-        "hashed_secret": "5b8b7a620e54e681c584f5b5c89152773c10c253",
-        "is_verified": false,
-        "line_number": 67
-      }
-    ],
     "core_backend/migrations/versions/2023_09_16_c5a948963236_create_query_table.py": [
       {
         "type": "Hex High Entropy String",
@@ -430,7 +421,7 @@
         "filename": "core_backend/tests/api/test_dashboard_overview.py",
         "hashed_secret": "6367c48dd193d56ea7b0baad25b19455e529f5ee",
         "is_verified": false,
-        "line_number": 291
+        "line_number": 290
       }
     ],
     "core_backend/tests/api/test_dashboard_performance.py": [
@@ -590,5 +581,5 @@
       }
     ]
   },
-  "generated_at": "2024-08-23T09:41:17Z"
+  "generated_at": "2024-08-26T14:03:01Z"
 }

core_backend/app/contents/schemas.py

@@ -1,22 +1,24 @@
 from datetime import datetime
-from typing import Annotated, List
+from typing import List
 
-from pydantic import BaseModel, ConfigDict, Field, StringConstraints
+from pydantic import BaseModel, ConfigDict, Field
 
 
 class ContentCreate(BaseModel):
     """
     Pydantic model for content creation request
     """
 
-    content_title: Annotated[str, StringConstraints(max_length=150)] = Field(
+    content_title: str = Field(
+        max_length=150,
         examples=["Example Content Title"],
     )
-    content_text: Annotated[str, StringConstraints(max_length=2000)] = Field(
-        examples=["This is an example content."]
+    content_text: str = Field(
+        max_length=2000,
+        examples=["This is an example content."],
     )
-    content_tags: list = Field(default=[], examples=[[1, 4]])
-    content_metadata: dict = Field(default={}, examples=[{"key": "optional_value"}])
+    content_tags: list = Field(default=[])
+    content_metadata: dict = Field(default={})
     is_archived: bool = False
 
     model_config = ConfigDict(

core_backend/app/llm_call/llm_rag.py

@@ -2,8 +2,6 @@
 Augmented Generation (RAG).
 """
 
-from typing import Optional
-
 from pydantic import ValidationError
 
 from ..config import LITELLM_MODEL_GENERATION
@@ -18,7 +16,7 @@ async def get_llm_rag_answer(
     question: str,
     context: str,
     original_language: IdentifiedLanguage,
-    metadata: Optional[dict] = None,
+    metadata: dict | None = None,
 ) -> RAG:
     """Get an answer from the LLM model using RAG.
 
@@ -32,7 +30,6 @@ async def get_llm_rag_answer(
         The language of the response.
     metadata
         Additional metadata to provide to the LLM model.
-
     Returns
     -------
     RAG

core_backend/app/llm_call/process_input.py

@@ -121,6 +121,7 @@ def _process_identified_language_response(
 
         error_response = QueryResponseError(
             query_id=response.query_id,
+            session_id=response.session_id,
             feedback_secret_key=response.feedback_secret_key,
             llm_response=response.llm_response,
             search_results=response.search_results,
@@ -206,6 +207,7 @@ async def _translate_question(
     else:
         error_response = QueryResponseError(
             query_id=response.query_id,
+            session_id=response.session_id,
             feedback_secret_key=response.feedback_secret_key,
             llm_response=response.llm_response,
             search_results=response.search_results,
@@ -275,6 +277,7 @@ async def _classify_safety(
     else:
         error_response = QueryResponseError(
             query_id=response.query_id,
+            session_id=response.session_id,
             feedback_secret_key=response.feedback_secret_key,
             llm_response=response.llm_response,
             search_results=response.search_results,
@@ -352,6 +355,7 @@ async def _paraphrase_question(
     else:
         error_response = QueryResponseError(
             query_id=response.query_id,
+            session_id=response.session_id,
             feedback_secret_key=response.feedback_secret_key,
             llm_response=response.llm_response,
             search_results=response.search_results,

core_backend/app/llm_call/process_output.py

@@ -45,39 +45,7 @@ class AlignScoreData(TypedDict):
     claim: str
 
 
-def generate_llm_response__after(func: Callable) -> Callable:
-    """
-    Decorator to generate the LLM response.
-
-    Only runs if the generate_llm_response flag is set to True.
-    Requires "search_results" and "original_language" in the response.
-    """
-
-    @wraps(func)
-    async def wrapper(
-        query_refined: QueryRefined,
-        response: QueryResponse | QueryResponseError,
-        *args: Any,
-        **kwargs: Any,
-    ) -> QueryResponse | QueryResponseError:
-        """
-        Generate the LLM response
-        """
-        response = await func(query_refined, response, *args, **kwargs)
-
-        if not query_refined.generate_llm_response:
-            return response
-
-        metadata = create_langfuse_metadata(
-            query_id=response.query_id, user_id=query_refined.user_id
-        )
-        response = await _generate_llm_response(query_refined, response, metadata)
-        return response
-
-    return wrapper
-
-
-async def _generate_llm_response(
+async def generate_llm_query_response(
     query_refined: QueryRefined,
     response: QueryResponse,
     metadata: Optional[dict] = None,
@@ -114,6 +82,7 @@ async def _generate_llm_response(
     else:
         response = QueryResponseError(
             query_id=response.query_id,
+            session_id=response.session_id,
             feedback_secret_key=response.feedback_secret_key,
             llm_response=None,
             search_results=response.search_results,
@@ -219,6 +188,7 @@ async def _check_align_score(
         )
         response = QueryResponseError(
             query_id=response.query_id,
+            session_id=response.session_id,
             feedback_secret_key=response.feedback_secret_key,
             llm_response=None,
             search_results=response.search_results,
@@ -311,6 +281,7 @@ async def wrapper(
             )
             response = QueryAudioResponse(
                 query_id=response.query_id,
+                session_id=response.session_id,
                 feedback_secret_key=response.feedback_secret_key,
                 llm_response=response.llm_response,
                 search_results=response.search_results,
@@ -361,6 +332,7 @@ async def _generate_tts_response(
         logger.error(f"Error generating TTS for query_id {response.query_id}: {e}")
         return QueryResponseError(
             query_id=response.query_id,
+            session_id=response.session_id,
             feedback_secret_key=response.feedback_secret_key,
             llm_response=response.llm_response,
             search_results=response.search_results,

core_backend/app/llm_call/utils.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from litellm import acompletion
 
 from ..config import LITELLM_API_KEY, LITELLM_ENDPOINT, LITELLM_MODEL_DEFAULT
@@ -11,9 +9,9 @@
 async def _ask_llm_async(
     user_message: str,
     system_message: str,
-    litellm_model: Optional[str] = LITELLM_MODEL_DEFAULT,
-    litellm_endpoint: Optional[str] = LITELLM_ENDPOINT,
-    metadata: Optional[dict] = None,
+    litellm_model: str | None = LITELLM_MODEL_DEFAULT,
+    litellm_endpoint: str | None = LITELLM_ENDPOINT,
+    metadata: dict | None = None,
     json: bool = False,
 ) -> str:
     """