Refactor agents and tools (#91)
* Move tools to agents * Move agents to dedicate place * Remove subclassing BaseAgent from BaseTool
This commit is contained in:
committed by
GitHub
parent
4256030b4f
commit
8e3a1d193f
34
knowledgehub/agents/tools/llm.py
Normal file
34
knowledgehub/agents/tools/llm.py
Normal file
@@ -0,0 +1,34 @@
|
||||
from typing import AnyStr, Optional, Type, Union
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from kotaemon.llms import LLM, AzureChatOpenAI, ChatLLM
|
||||
|
||||
from .base import BaseTool, ToolException
|
||||
|
||||
BaseLLM = Union[ChatLLM, LLM]
|
||||
|
||||
|
||||
class LLMArgs(BaseModel):
|
||||
query: str = Field(..., description="a search question or prompt")
|
||||
|
||||
|
||||
class LLMTool(BaseTool):
|
||||
name = "llm"
|
||||
description = (
|
||||
"A pretrained LLM like yourself. Useful when you need to act with "
|
||||
"general world knowledge and common sense. Prioritize it when you "
|
||||
"are confident in solving the problem "
|
||||
"yourself. Input can be any instruction."
|
||||
)
|
||||
llm: BaseLLM = AzureChatOpenAI()
|
||||
args_schema: Optional[Type[BaseModel]] = LLMArgs
|
||||
|
||||
def _run_tool(self, query: AnyStr) -> str:
|
||||
output = None
|
||||
try:
|
||||
response = self.llm(query)
|
||||
except ValueError:
|
||||
raise ToolException("LLM Tool call failed")
|
||||
output = response.text
|
||||
return output
|
Reference in New Issue
Block a user