3cceec63ef
* add base Tool * minor update test_tool * update test dependency * update test dependency * Fix namespace conflict * update test * add base Agent Interface, add ReWoo Agent * minor update * update test * fix typo * remove unneeded print * update rewoo agent * add LLMTool * update BaseAgent type * add ReAct agent * add ReAct agent * minor update * minor update * minor update * minor update * update docstring * fix max_iteration --------- Co-authored-by: trducng <trungduc1992@gmail.com>
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
from typing import AnyStr, Optional, Type, Union
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from kotaemon.llms.chats.base import ChatLLM
|
|
from kotaemon.llms.chats.openai import AzureChatOpenAI
|
|
from kotaemon.llms.completions.base import LLM
|
|
|
|
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[0]
|
|
return output
|