fix typing

onyx-dot-app · Sep 2, 2024 · 16f9405 · 16f9405
1 parent 9e5149e
commit 16f9405
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 33 deletions.
diff --git a/backend/danswer/one_shot_answer/answer_question.py b/backend/danswer/one_shot_answer/answer_question.py
@@ -83,7 +83,6 @@
 def stream_answer_objects(
     query_req: DirectQARequest,
     user: User | None,
-    temporary_persona: Persona | None,
     # These need to be passed in because in Web UI one shot flow,
     # we can have much more document as there is no history.
     # For Slack flow, we need to save more tokens for the thread context
@@ -99,6 +98,7 @@ def stream_answer_objects(
     retrieval_metrics_callback: (
         Callable[[RetrievalMetricsContainer], None] | None
     ) = None,
+    temporary_persona: Persona | None = None,
     rerank_metrics_callback: Callable[[RerankMetricsContainer], None] | None = None,
 ) -> AnswerObjectIterator:
     """Streams in order:
@@ -184,33 +184,27 @@ def stream_answer_objects(
         max_tokens=max_document_tokens,
     )
 
-    contains_tool = True
     if temporary_persona:
-        contains_tool = False
         for tool in temporary_persona.tools:
             if tool.in_code_tool_id == "SearchTool":
-                contains_tool = True
+                pass
 
-    search_tool = (
-        SearchTool(
-            db_session=db_session,
-            user=user,
-            evaluation_type=LLMEvaluationType.SKIP
-            if DISABLE_LLM_DOC_RELEVANCE
-            else query_req.evaluation_type,
-            persona=persona,
-            retrieval_options=query_req.retrieval_options,
-            prompt_config=prompt_config,
-            llm=llm,
-            fast_llm=fast_llm,
-            pruning_config=document_pruning_config,
-            bypass_acl=bypass_acl,
-            chunks_above=query_req.chunks_above,
-            chunks_below=query_req.chunks_below,
-            full_doc=query_req.full_doc,
-        )
-        if contains_tool
-        else None
+    search_tool = SearchTool(
+        db_session=db_session,
+        user=user,
+        evaluation_type=LLMEvaluationType.SKIP
+        if DISABLE_LLM_DOC_RELEVANCE
+        else query_req.evaluation_type,
+        persona=persona,
+        retrieval_options=query_req.retrieval_options,
+        prompt_config=prompt_config,
+        llm=llm,
+        fast_llm=fast_llm,
+        pruning_config=document_pruning_config,
+        bypass_acl=bypass_acl,
+        chunks_above=query_req.chunks_above,
+        chunks_below=query_req.chunks_below,
+        full_doc=query_req.full_doc,
     )
 
     answer_config = AnswerStyleConfig(
@@ -230,9 +224,8 @@ def stream_answer_objects(
             ForceUseTool(
                 tool_name=search_tool.name,
                 args={"query": rephrased_query},
+                force_use=True,
             )
-            if search_tool
-            else None
         ),
         # for now, don't use tool calling for this flow, as we haven't
         # tested quotes with tool calling too much yet

diff --git a/backend/danswer/one_shot_answer/models.py b/backend/danswer/one_shot_answer/models.py
@@ -1,3 +1,5 @@
+from typing import Any
+
 from pydantic import BaseModel
 from pydantic import Field
 from pydantic import model_validator
@@ -8,9 +10,9 @@
 from danswer.chat.models import QADocsResponse
 from danswer.configs.constants import MessageType
 from danswer.db.models import StarterMessage
+from danswer.search.enums import LLMEvaluationType
 from danswer.search.enums import RecencyBiasSetting
 from danswer.search.enums import SearchType
-from danswer.search.enums import LLMEvaluationType
 from danswer.search.models import ChunkContext
 from danswer.search.models import RerankingDetails
 from danswer.search.models import RetrievalDetails
@@ -40,13 +42,13 @@ class DocumentSetConfig(BaseModel):
 
 
 class ToolConfig(BaseModel):
-    id: int | None = None
+    id: int
 
 
 class PersonaConfig(BaseModel):
     name: str
     description: str
-    search_type: SearchType = SearchType.HYBRID
+    search_type: SearchType = SearchType.SEMANTIC
     num_chunks: float | None = None
     llm_relevance_filter: bool = False
     llm_filter_extraction: bool = False
@@ -67,11 +69,10 @@ class PersonaConfig(BaseModel):
 
 class DirectQARequest(ChunkContext):
     persona_config: PersonaConfig | None = None
-    persona_id: int | None = None
+    persona_id: int
 
     messages: list[ThreadMessage]
     prompt_id: int | None
-    persona_id: int
     multilingual_query_expansion: list[str] | None = None
     retrieval_options: RetrievalDetails = Field(default_factory=RetrievalDetails)
     rerank_settings: RerankingDetails | None = None

diff --git a/backend/danswer/server/query_and_chat/models.py b/backend/danswer/server/query_and_chat/models.py
@@ -136,7 +136,7 @@ class RenameChatSessionResponse(BaseModel):
 class ChatSessionDetails(BaseModel):
     id: int
     name: str
-    persona_id: int
+    persona_id: int | None = None
     time_created: str
     shared_status: ChatSessionSharedStatus
     folder_id: int | None = None
@@ -196,7 +196,7 @@ class SearchSessionDetailResponse(BaseModel):
 class ChatSessionDetailResponse(BaseModel):
     chat_session_id: int
     description: str
-    persona_id: int
+    persona_id: int | None = None
     persona_name: str
     messages: list[ChatMessageDetail]
     time_created: datetime