BUG: Fix api endpoint handlers. (#214)

web/api/chat_completion_handler.py

@@ -1,6 +1,7 @@
 """Handle /api/chat/completion requests."""
-from typing import Optional, Self
+from typing import Any, Optional, Self
 
+from docq.manage_personas import get_persona
 from docq.model_selection.main import get_model_settings_collection
 from docq.run_queries import run_chat
 from pydantic import Field, ValidationError
@@ -15,6 +16,7 @@ class PostRequestModel(CamelModel):
     input_: str = Field(..., alias="input")
     history: Optional[str] = Field(None)
     llm_settings_collection_name: Optional[str] = Field(None)
+    persona_key: Optional[str] = Field(None)
 
 class PostResponseModel(CamelModel):
     """Pydantic model for the response body."""
@@ -25,11 +27,12 @@ class PostResponseModel(CamelModel):
 class ChatCompletionHandler(RequestHandler):
     """Handle /api/chat/completion requests."""
 
-    def check_origin(self: Self, origin) -> bool:
+    def check_origin(self: Self, origin: Any) -> bool:
         """Override the origin check if it's causing problems."""
         return True
 
-    def check_xsrf_cookie(self) -> bool:
+    def check_xsrf_cookie(self: Self) -> bool:
+        """Override the XSRF cookie check."""
         # If `True`, POST, PUT, and DELETE are block unless the `_xsrf` cookie is set.
         # Safe with token based authN
         return False
@@ -56,7 +59,8 @@ def post(self: Self) -> None:
             request_model = PostRequestModel.model_validate_json(body)
             history = request_model.history if request_model.history else ""
             model_usage_settings = get_model_settings_collection(request_model.llm_settings_collection_name) if request_model.llm_settings_collection_name else get_model_settings_collection("azure_openai_latest")
-            result = run_chat(input_=request_model.input_, history=history, model_settings_collection=model_usage_settings)
+            persona = get_persona(request_model.persona_key if request_model.persona_key else "default")
+            result = run_chat(input_=request_model.input_, history=history, model_settings_collection=model_usage_settings, persona=persona)
             response_model = PostResponseModel(response=result.response, meta={"model_settings": model_usage_settings.key})
 
             self.write(response_model.model_dump())

web/api/utils.py

@@ -4,25 +4,24 @@
 import re
 from typing import Any, Callable, Iterable, Mapping
 
+from docq.config import ENV_VAR_DOCQ_API_SECRET
 from opentelemetry import trace
 from pydantic import BaseModel
 from tornado.web import HTTPError, RequestHandler
 
-from ...source.docq.config import ENV_VAR_DOCQ_API_SECRET
-
 tracer = trace.get_tracer(__name__)
 
 UNDERSCORE_RE = re.compile(r"(?<=[^\-_])[\-_]+[^\-_]")
 
-def _process_keys(str_or_iter: str | Iterable, fn) -> str | Iterable:
+def _process_keys(str_or_iter: str | Iterable, fn: Any) -> str | Iterable:
     """Recursively process keys in a string, dict, or list of dicts."""
     if isinstance(str_or_iter, list):
         return [_process_keys(k, fn) for k in str_or_iter]
     if isinstance(str_or_iter, Mapping):
         return {fn(k): _process_keys(v, fn) for k, v in str_or_iter.items()}
     return str_or_iter
 
-def _is_none(_in) -> str:
+def _is_none(_in: Any) -> str:
     """Determine if the input is None and returns a string with white-space removed.
 
     :param _in: input.
@@ -60,14 +59,15 @@ def camelize(str_or_iter: str | Iterable) -> str | Iterable:
     return UNDERSCORE_RE.sub(lambda m: m.group(0)[-1].upper(), s)
 
 
-def to_camel(string) -> str | Iterable:
+def to_camel(string: str) -> str | Iterable:
     """Convert a string to camel case."""
     return camelize(string)
 
 
 class CamelModel(BaseModel):
     """Pydantic model that generated camelCase alias from snake_case field names."""
     class Config:
+        """Pydantic model configuration."""
         alias_generator = to_camel
         population_by_name = True
 

web/index.py

@@ -15,6 +15,8 @@
 )
 from utils.observability import baggage_as_attributes, tracer
 
+import web.api.index_handler  # noqa F401 don't remove this line, it's used to register api routes
+
 with tracer().start_as_current_span("home_page", attributes=baggage_as_attributes()):
     render_page_title_and_favicon()
     init_with_pretty_error_ui()