highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-06-16 19:10:53 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix typing for custom tool response

Browse Source
This commit is contained in:
Weves 2024-08-23 11:13:36 -07:00 committed by Chris Weaver
parent e749fa0f28
commit 7d201f67d4
3 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
backend/danswer/chat/models.py
View File

@ -9,6 +9,7 @@ from danswer.search.enums import QueryFlow
from danswer.search.enums import SearchType from danswer.search.enums import SearchType
from danswer.search.models import RetrievalDocs from danswer.search.models import RetrievalDocs
from danswer.search.models import SearchResponse from danswer.search.models import SearchResponse
from danswer.tools.custom.base_tool_types import ToolResultType
class LlmDoc(BaseModel): class LlmDoc(BaseModel):
@ -130,7 +131,7 @@ class ImageGenerationDisplay(BaseModel):
class CustomToolResponse(BaseModel): class CustomToolResponse(BaseModel):
response: dict response: ToolResultType
tool_name: str tool_name: str

2
backend/danswer/tools/custom/base_tool_types.py Normal file
View File

@ -0,0 +1,2 @@
# should really be `JSON_ro`, but this causes issues with pydantic
ToolResultType = dict | list | str | int | float | bool

3
backend/danswer/tools/custom/custom_tool.py
View File

@ -11,6 +11,7 @@ from pydantic import BaseModel
from danswer.dynamic_configs.interface import JSON_ro from danswer.dynamic_configs.interface import JSON_ro
from danswer.llm.answering.models import PreviousMessage from danswer.llm.answering.models import PreviousMessage
from danswer.llm.interfaces import LLM from danswer.llm.interfaces import LLM
from danswer.tools.custom.base_tool_types import ToolResultType
from danswer.tools.custom.custom_tool_prompts import ( from danswer.tools.custom.custom_tool_prompts import (
SHOULD_USE_CUSTOM_TOOL_SYSTEM_PROMPT, SHOULD_USE_CUSTOM_TOOL_SYSTEM_PROMPT,
) )
@ -34,7 +35,7 @@ CUSTOM_TOOL_RESPONSE_ID = "custom_tool_response"
class CustomToolCallSummary(BaseModel): class CustomToolCallSummary(BaseModel):
tool_name: str tool_name: str
tool_result: dict tool_result: ToolResultType
class CustomTool(Tool): class CustomTool(Tool):