book_rag_server/utils/ragflow/agent_service.py

from typing import Dict, Any, List, Optional

class AgentService:
    def __init__(self, http_client):
        self.http_client = http_client
    
    def create_agent(self, name: str, llm: Dict[str, Any], description: str = None) -> Dict[str, Any]:
        endpoint = "/api/v1/agents"
        
        data = {"name": name, "llm": llm}
        if description is not None:
            data["description"] = description
        
        response = self.http_client.post(endpoint, json_data=data)
        
        if response.get("code") == 0 and response.get("data"):
            return response["data"]
        else:
            raise Exception(f"创建代理失败: {response.get('message', '未知错误')}")
    
    def update_agent(self, agent_id: str, name: str = None, llm: Dict[str, Any] = None,
                    description: str = None) -> Dict[str, Any]:
        endpoint = f"/api/v1/agents/{agent_id}"
        
        data = {}
        if name is not None:
            data["name"] = name
        if llm is not None:
            data["llm"] = llm
        if description is not None:
            data["description"] = description
        
        response = self.http_client.post(endpoint, json=data)
        
        if response.get("code") == 0 and response.get("data"):
            return response["data"]
        else:
            raise Exception(f"更新代理失败: {response.get('message', '未知错误')}")
    
    def delete_agent(self, agent_id: str) -> bool:
        endpoint = f"/api/v1/agents/{agent_id}"
        
        response = self.http_client.post(endpoint, json_data={})
        
        if response.get("code") == 0:
            return True
        else:
            raise Exception(f"删除代理失败: {response.get('message', '未知错误')}")
    
    def list_agents(self, page: int = 1, size: int = 20, orderby: str = "create_time",
                   desc: bool = True, name: str = None, agent_id: str = None) -> List[Dict[str, Any]]:
        endpoint = "/api/v1/agents"
        
        params = {"page": page, "page_size": size, "orderby": orderby, "desc": int(desc)}
        if name is not None:
            params["name"] = name
        if agent_id is not None:
            params["id"] = agent_id
        
        response = self.http_client.get(endpoint, params=params)
        
        if response.get("code") == 0 and response.get("data"):
            return response["data"]
        else:
            raise Exception(f"列出代理失败: {response.get('message', '未知错误')}")
    
    def create_agent_session(self, agent_id: str, name: str = None) -> Dict[str, Any]:
        endpoint = f"/api/v1/agents/{agent_id}/sessions"
        
        data = {}
        if name is not None:
            data["name"] = name
        
        response = self.http_client.post(endpoint, json=data)
        
        if response.get("code") == 0 and response.get("data"):
            return response["data"]
        else:
            raise Exception(f"创建代理会话失败: {response.get('message', '未知错误')}")
    
    def list_agent_sessions(self, agent_id: str, page: int = 1, size: int = 20,
                           orderby: str = "create_time", desc: bool = True,
                           session_id: str = None, user_id: str = None,
                           dsl: str = None) -> List[Dict[str, Any]]:
        endpoint = f"/api/v1/agents/{agent_id}/sessions"
        
        params = {"page": page, "page_size": size, "orderby": orderby, "desc": int(desc)}
        if session_id is not None:
            params["id"] = session_id
        if user_id is not None:
            params["user_id"] = user_id
        if dsl is not None:
            params["dsl"] = dsl
        
        response = self.http_client.get(endpoint, params=params)
        
        if response.get("code") == 0 and response.get("data"):
            return response["data"]
        else:
            raise Exception(f"列出代理会话失败: {response.get('message', '未知错误')}")
    
    def delete_agent_session(self, agent_id: str, session_id: str) -> bool:
        endpoint = f"/api/v1/agents/{agent_id}/sessions"
        
        response = self.http_client.post(endpoint, json_data={"session_ids": [session_id]})
        
        if response.get("code") == 0:
            return True
        else:
            raise Exception(f"删除代理会话失败: {response.get('message', '未知错误')}")
    
    def agent_completion(self, agent_id: str, query: str, stream: bool = False,
                        session_id: str = None) -> Dict[str, Any]:
        endpoint = f"/api/v1/agents/{agent_id}/completions"
        
        data = {"query": query, "stream": stream}
        if session_id is not None:
            data["session_id"] = session_id
        
        response = self.http_client.post(endpoint, json=data)
        
        if response.get("code") == 0:
            return response.get("data", {})
        else:
            raise Exception(f"代理完成失败: {response.get('message', '未知错误')}")
    
    def get_related_questions(self, dataset_id: str, question: str, top: int = 10) -> List[str]:
        endpoint = "/api/v1/sessions/related_questions"
        
        response = self.http_client.post(endpoint, json_data={
            "dataset_id": dataset_id,
            "question": question,
            "top": top
        })
        
        if response.get("code") == 0 and response.get("data"):
            return response["data"]
        else:
            raise Exception(f"获取相关问题失败: {response.get('message', '未知错误')}")