RagFlow_Plugs/services/ragflow/chat_service.py

from typing import Dict, Any, List, Optional

class ChatService:
    def __init__(self, http_client):
        self.http_client = http_client
    
    def create_chat(self, name: str, dataset_ids: List[str], llm: Dict[str, Any],
                   prompt: str = None) -> Dict[str, Any]:
        endpoint = "/api/v1/chats"
        
        data = {
            "name": name,
            "dataset_ids": dataset_ids,
            "llm": llm
        }
        if prompt is not None:
            data["prompt"] = prompt
        
        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 update_chat(self, chat_id: str, name: str = None, dataset_ids: List[str] = None,
                   llm: Dict[str, Any] = None, prompt: str = None) -> Dict[str, Any]:
        endpoint = f"/api/v1/chats/{chat_id}"
        
        data = {}
        if name is not None:
            data["name"] = name
        if dataset_ids is not None:
            data["dataset_ids"] = dataset_ids
        if llm is not None:
            data["llm"] = llm
        if prompt is not None:
            data["prompt"] = prompt
        
        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_chats(self, chat_ids: List[str]) -> bool:
        endpoint = "/api/v1/chats"
        
        response = self.http_client.post(endpoint, json={"chat_ids": chat_ids})
        
        if response.get("code") == 0:
            return True
        else:
            raise Exception(f"删除聊天失败: {response.get('message', '未知错误')}")
    
    def list_chats(self, page: int = 1, size: int = 20, orderby: str = "create_time",
                  desc: bool = True, name: str = None, chat_id: str = None) -> List[Dict[str, Any]]:
        endpoint = "/api/v1/chats"
        
        params = {"page": page, "page_size": size, "orderby": orderby, "desc": int(desc)}
        if name is not None:
            params["name"] = name
        if chat_id is not None:
            params["id"] = chat_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_chat_session(self, chat_id: str, name: str = None) -> Dict[str, Any]:
        endpoint = f"/api/v1/chats/{chat_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 update_chat_session(self, chat_id: str, session_id: str, 
                           name: str = None, message: List[Dict] = None) -> Dict[str, Any]:
        endpoint = f"/api/v1/chats/{chat_id}/sessions/{session_id}"
        
        data = {}
        if name is not None:
            data["name"] = name
        if message is not None:
            data["message"] = message
        
        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_chat_sessions(self, chat_id: str, page: int = 1, size: int = 20,
                          orderby: str = "create_time", desc: bool = True,
                          session_id: str = None, session_name: str = None) -> List[Dict[str, Any]]:
        endpoint = f"/api/v1/chats/{chat_id}/sessions"
        
        params = {"page": page, "page_size": size, "orderby": orderby, "desc": int(desc)}
        if session_id is not None:
            params["id"] = session_id
        if session_name is not None:
            params["name"] = session_name
        
        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_chat_session(self, chat_id: str, session_id: str) -> bool:
        endpoint = f"/api/v1/chats/{chat_id}/sessions"
        
        response = self.http_client.post(endpoint, json={"session_ids": [session_id]})
        
        if response.get("code") == 0:
            return True
        else:
            raise Exception(f"删除会话失败: {response.get('message', '未知错误')}")
    
    def chat_completion(self, chat_id: str, query: str, stream: bool = False,
                       session_id: str = None) -> Dict[str, Any]:
        endpoint = f"/api/v1/chats/{chat_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', '未知错误')}")