Merge pull request #2 from difizen/brokun_dev

refactor(API): core agent

api/core/agent.py

@@ -0,0 +1,32 @@
+'''Agent'''
+from typing import AsyncIterator, List
+from abc import ABC, abstractmethod
+from langchain.schema.messages import BaseMessage, BaseMessageChunk
+from pydantic import BaseModel
+
+from models.chat import MessageModel
+
+from .base import ConfigMeta
+
+
+class Agent(BaseModel, ABC):
+    meta: ConfigMeta
+
+    '''Abstract agent'''
+    @abstractmethod
+    def invoke(
+        self,
+        config: ConfigMeta,
+        chat_id: int,
+        history: List[MessageModel],
+        message: MessageModel,
+    ) -> BaseMessage | None:
+        """Chat with agent"""
+
+    @abstractmethod
+    def invoke_astream(self,
+                       config: ConfigMeta,
+                       chat_id: int,
+                       history: List[MessageModel],
+                       message: MessageModel,) -> AsyncIterator[BaseMessageChunk] | None:
+        """Chat with agent and get answer via stream"""

api/core/agent_manager.py

@@ -0,0 +1,40 @@
+'''Agent manager
+'''
+
+
+from typing import List
+from pydantic import BaseModel
+
+from .agent_provider import AgentProvider
+from .base import ConfigMeta
+
+
+class AgentManager(BaseModel):
+    '''Agent Manager
+    get or create agent from provider by manager
+    '''
+
+    providers: List[AgentProvider] = []
+
+    def registe_provider(self, provider: AgentProvider):
+        '''add provider'''
+        self.providers.append(provider)
+
+    def get_provider(self, config: ConfigMeta) -> AgentProvider:
+        '''retrieve the best matching provider'''
+        selected_provider = None
+        highest_priority = -1
+
+        for provider in self.providers:
+            priority = provider.can_handle(config)
+            if priority > highest_priority:
+                highest_priority = priority
+                selected_provider = provider
+
+        if selected_provider is None:
+            raise Exception("No suitable provider found")
+
+        return selected_provider
+
+
+agent_manager = AgentManager()

api/core/agent_provider.py

@@ -0,0 +1,20 @@
+'''Agent Provider'''
+from abc import ABC, abstractmethod
+from pydantic import BaseModel
+
+from .agent import Agent
+
+from .base import ConfigMeta
+
+
+class AgentProvider(BaseModel, ABC):
+    '''Agent Provider'''
+    name: str = "custom"
+
+    def can_handle(self, config: ConfigMeta) -> int:
+        '''Priority of provider for handling agent config'''
+        return -1
+
+    @abstractmethod
+    def provide(self, config: ConfigMeta, chat_id: int) -> Agent:
+        """List chat executors."""

api/core/base.py

@@ -0,0 +1,15 @@
+'''Base model'''
+from typing import Any, List
+from pydantic import BaseModel
+
+
+class ModelMeta(BaseModel):
+    '''LLM meta model '''
+    key: str
+
+
+class ConfigMeta(BaseModel):
+    '''agent meta model'''
+    persona: str
+    model: ModelMeta
+    tools: List[Any] = []

api/core/chat.py

@@ -1,19 +1,14 @@
 from typing import List
 
 from langchain.schema.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage
-from pydantic import BaseModel
 from models.agent_config import AgentConfigModel
 from models.chat import MessageModel, MessageSenderType
-from .chat_executor import chat_object_manager
+from .agent_manager import agent_manager
+from .base import ConfigMeta
 
+from .langchain_agent_provider import langchain_agent_provider
 
-class ModelMeta(BaseModel):
-    key: str
-
-
-class ConfigMeta(BaseModel):
-    persona: str
-    model: ModelMeta
+agent_manager.registe_provider(langchain_agent_provider)
 
 
 def get_config_meta(agent_config: AgentConfigModel) -> ConfigMeta:
@@ -35,8 +30,9 @@ def chat(agent_config: AgentConfigModel, chat_id: int, history: List[MessageMode
     system_msg = SystemMessage(content=config.persona)
     msgs = [to_message(m) for m in history]
     msgs.insert(0, system_msg)
-    executor = chat_object_manager.get_executor(config.model.key)
-    answer = executor.run(msgs)
+    provider = agent_manager.get_provider(config)
+    agent = provider.provide(config, chat_id)
+    answer = agent.invoke(config, chat_id, history, message)
     return answer
 
 
@@ -45,6 +41,7 @@ def chat_stream(agent_config: AgentConfigModel, chat_id: int, history: List[Mess
     system_msg = SystemMessage(content=config.persona)
     msgs = [to_message(m) for m in history]
     msgs.insert(0, system_msg)
-    executor = chat_object_manager.get_executor(config.model.key)
-    answer = executor.astream(msgs)
+    provider = agent_manager.get_provider(config)
+    agent = provider.provide(config, chat_id)
+    answer = agent.invoke_astream(config, chat_id, history, message)
     return answer