shudao-main/shudao-chat-py/services/chromadb_service.py

"""
ChromaDB 知识库检索服务
"""
import httpx
from typing import List, Dict, Any
from utils.logger import logger


class ChromaDBService:
    def __init__(self):
        # ChromaDB 配置可以从 config 中读取，这里暂时硬编码
        self.base_url = "http://localhost:8000"  # ChromaDB 服务地址
    
    async def query_documents(self, query: str, n_results: int = 5) -> List[Dict[str, Any]]:
        """查询相关文档"""
        try:
            data = {
                "query": query,
                "n_results": n_results
            }
            
            async with httpx.AsyncClient(timeout=10.0) as client:
                response = await client.post(
                    f"{self.base_url}/query",
                    json=data
                )
                response.raise_for_status()
                return response.json().get('documents', [])
        except httpx.HTTPError as e:
            logger.warning(f"ChromaDB API 调用失败，返回模拟数据: {e}")
            # 返回模拟数据
            return [
                {
                    "content": f"关于 {query} 的相关文档内容...",
                    "metadata": {"source": "模拟数据"},
                    "distance": 0.5
                }
            ]
        except Exception as e:
            logger.error(f"ChromaDB 服务异常: {e}")
            return []


# 全局实例
chromadb_service = ChromaDBService()