kotaemon/knowledgehub/pipelines/agents/base.py
Nguyen Trung Duc (john) 693ed39de4 Move prompts into LLMs module (#70)
Since the only usage of prompt is within LLMs, it is reasonable to keep it within the LLM module. This way, it would be easier to discover module, and make the code base less complicated.

Changes:

* Move prompt components into llms
* Bump version 0.3.1
* Make pip install dependencies in eager mode

---------

Co-authored-by: ian <ian@cinnamon.is>
2023-11-14 16:00:10 +07:00

68 lines
1.8 KiB
Python

from enum import Enum
from typing import Dict, List, Optional, Union
from pydantic import BaseModel
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"
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."""
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
"""