8bb7ad91e0
* update Param() type hint in MVP * update default embedding endpoint * update Langchain agent wrapper * update langchain agent
62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
from enum import Enum
|
|
from typing import Dict, List, Optional, Union
|
|
|
|
from kotaemon.llms import PromptTemplate
|
|
from kotaemon.llms.chats.base import ChatLLM
|
|
from kotaemon.llms.completions.base import LLM
|
|
from kotaemon.pipelines.tools import BaseTool
|
|
|
|
BaseLLM = Union[ChatLLM, LLM]
|
|
|
|
|
|
class AgentType(Enum):
|
|
"""
|
|
Enumerated type for agent types.
|
|
"""
|
|
|
|
openai = "openai"
|
|
openai_multi = "openai_multi"
|
|
openai_tool = "openai_tool"
|
|
self_ask = "self_ask"
|
|
react = "react"
|
|
rewoo = "rewoo"
|
|
vanilla = "vanilla"
|
|
|
|
@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 BaseAgent(BaseTool):
|
|
name: str
|
|
"""Name of the agent."""
|
|
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: Optional[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
|
|
"""
|
|
|
|
def add_tools(self, tools: List[BaseTool]) -> None:
|
|
"""Helper method to add tools and update agent state if needed"""
|
|
self.plugins.extend(tools)
|