highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-05-29 17:19:36 +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.models import RetrievalDocs
from danswer.search.models import SearchResponse
from danswer.tools.custom.base_tool_types import ToolResultType
class LlmDoc(BaseModel):
@ -130,7 +131,7 @@ class ImageGenerationDisplay(BaseModel):
class CustomToolResponse(BaseModel):
response: dict
response: ToolResultType
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.llm.answering.models import PreviousMessage
from danswer.llm.interfaces import LLM
from danswer.tools.custom.base_tool_types import ToolResultType
from danswer.tools.custom.custom_tool_prompts import (
SHOULD_USE_CUSTOM_TOOL_SYSTEM_PROMPT,
)
@ -34,7 +35,7 @@ CUSTOM_TOOL_RESPONSE_ID = "custom_tool_response"
class CustomToolCallSummary(BaseModel):
tool_name: str
tool_result: dict
tool_result: ToolResultType
class CustomTool(Tool):