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>
68 lines
1.8 KiB
Python
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
|
|
"""
|