91048770fa
* 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 --------- Co-authored-by: trducng <trungduc1992@gmail.com>
68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
from enum import Enum
|
|
from typing import Dict, List, Union
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from kotaemon.llms.chats.base import ChatLLM
|
|
from kotaemon.llms.completions.base import LLM
|
|
from kotaemon.pipelines.tools import BaseTool
|
|
from kotaemon.prompt.template import PromptTemplate
|
|
|
|
BaseLLM = Union[ChatLLM, LLM]
|
|
|
|
|
|
class AgentType(Enum):
|
|
"""
|
|
Enumerated type for agent types.
|
|
"""
|
|
|
|
openai = "openai"
|
|
react = "react"
|
|
rewoo = "rewoo"
|
|
vanilla = "vanilla"
|
|
openai_memory = "openai_memory"
|
|
|
|
@staticmethod
|
|
def get_agent_class(_type: "AgentType"):
|
|
"""
|
|
Get agent class from agent type.
|
|
:param _type: agent type
|
|
:return: agent class
|
|
"""
|
|
if _type == AgentType.rewoo:
|
|
from .rewoo.agent import RewooAgent
|
|
|
|
return RewooAgent
|
|
else:
|
|
raise ValueError(f"Unknown agent type: {_type}")
|
|
|
|
|
|
class AgentOutput(BaseModel):
|
|
"""
|
|
Pydantic model for agent output.
|
|
"""
|
|
|
|
output: str
|
|
cost: float
|
|
token_usage: int
|
|
|
|
|
|
class BaseAgent(BaseTool):
|
|
name: str
|
|
"""Name of the agent."""
|
|
type: AgentType
|
|
"""Agent type, must be one of AgentType"""
|
|
description: str
|
|
"""Description used to tell the model how/when/why to use the agent.
|
|
You can provide few-shot examples as a part of the description. This will be
|
|
input to the prompt of LLM."""
|
|
llm: Union[BaseLLM, Dict[str, BaseLLM]]
|
|
"""Specify LLM to be used in the model, cam be a dict to supply different
|
|
LLMs to multiple purposes in the agent"""
|
|
prompt_template: Union[PromptTemplate, Dict[str, PromptTemplate]]
|
|
"""A prompt template or a dict to supply different prompt to the agent
|
|
"""
|
|
plugins: List[BaseTool]
|
|
"""List of plugins / tools to be used in the agent
|
|
"""
|