diff --git a/docker-compose.yml b/docker-compose.yml index 6afa178..e2ed891 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,17 +1,29 @@ +networks: + ai-pipeline: + driver: bridge + +volumes: + es-data: + qdrant-storage: + services: elasticsearch: - image: elasticsearch:9.2.2 + image: docker.elastic.co/elasticsearch/elasticsearch:9.2.2 container_name: elasticsearch + user: "1000:1000" environment: - discovery.type=single-node # 单节点模式 - - xpack.security.enabled=true # 开启安全(才可以设密码) - - ELASTIC_PASSWORD=12345 # 设置密码(重要) - - ES_JAVA_OPTS=-Xms1g -Xmx1g + - xpack.security.enabled=true # 开启安全 + - ELASTIC_PASSWORD=12345 # 设置密码 + # - xpack.security.http.ssl.enabled=false # 关闭 HTTPS + # - ES_JAVA_OPTS=-Xms1g -Xmx1g # JVM 初始堆内存大小,1GB,最大堆内存大小,1GB ports: - "9210:9200" volumes: - - ../datas/es-data:/usr/share/elasticsearch/data + - es-data:/usr/share/elasticsearch/data restart: unless-stopped + networks: + - ai-pipeline qdrant: image: qdrant/qdrant:latest @@ -22,5 +34,7 @@ services: - "6333:6333" # HTTP - "6334:6334" # gRPC volumes: - - ../datas/qdrant-storage:/qdrant/storage - restart: unless-stopped \ No newline at end of file + - qdrant-storage:/qdrant/storage + restart: unless-stopped + networks: + - ai-pipeline diff --git a/src/pipeline/config.py b/src/pipeline/config.py index 71d4aa2..2f3130b 100644 --- a/src/pipeline/config.py +++ b/src/pipeline/config.py @@ -1,6 +1,7 @@ from mimetypes import init from typing import TypedDict from dotenv import load_dotenv +from src.pipeline.core.utils import logger import os load_dotenv() @@ -17,9 +18,12 @@ class Config(TypedDict): embedding_api_host: str embedding_model: str es_host: str - es_key: str + es_port: int + es_user: str + es_password: str qdrant_host: str - qdrant_key: str + qdrant_port: int + qdrant_api_key: str def _read_config() -> Config: @@ -35,9 +39,13 @@ def _read_config() -> Config: "embedding_api_key": os.getenv("EMBEDDING_API_KEY"), "embedding_model": os.getenv("EMBEDDING_MODEL"), "es_host": os.getenv("ES_HOST"), - "es_key": os.getenv("ES_KEY") or "", + "es_port": int(os.getenv("ES_PORT")), + "es_host": os.getenv("ES_USER") or "elastic", + "es_password": os.getenv("ES_PASSWORD") or "", "qdrant_host": os.getenv("QDRANT_HOST"), - "qdrant_host": os.getenv("QDRANT_KEY") or "", + "qdrant_port": int(os.getenv("QDRANT_PORT")), + "qdrant_api_key": os.getenv("QDRANT_API_KEY") or "", } config = _read_config() +logger.debug("创建全局: config") diff --git a/src/pipeline/core/es_client.py b/src/pipeline/core/es_client.py new file mode 100644 index 0000000..1718597 --- /dev/null +++ b/src/pipeline/core/es_client.py @@ -0,0 +1,45 @@ +import httpx +from pipeline.config import config + +class AsyncES: + def __init__(self): + self.base = f"{config['es_host']}:{config['es_port']}" + self.auth = (config["es_user"], config['es_password']) + + async def create_index(self, index: str): + """ + 创建 index(知识库) + """ + async with httpx.AsyncClient() as client: + resp = await client.put( + f"{self.base}/{index}", + auth=self.auth, + json={ + "settings": {"analysis": {"analyzer": {"default": {"type": "standard"}}}}, + "mappings": { + "properties": { + "text": {"type": "text"}, + "kb": {"type": "keyword"}, + } + }, + }, + ) + return resp.json() + + async def add_doc(self, index: str, text: str, kb: str): + """ + 写入文档(普通文本) + """ + async with httpx.AsyncClient() as client: + resp = await client.post(f"{self.base}/{index}/_doc", auth=self.auth, json={"text": text, "kb": kb}) + return resp.json() + + async def bm25_search(self, index: str, query: str, top_k: int = 10): + """ + BM25 搜索 + """ + async with httpx.AsyncClient() as client: + resp = await client.post(f"{self.base}/{index}/_search", auth=self.auth, json={"size": top_k, "query": {"match": {"text": query}}}) + return resp.json() + +es_client = AsyncES diff --git a/src/pipeline/core/qd_client.py b/src/pipeline/core/qd_client.py new file mode 100644 index 0000000..413280a --- /dev/null +++ b/src/pipeline/core/qd_client.py @@ -0,0 +1,53 @@ +from qdrant_client import AsyncQdrantClient +from qdrant_client.models import ( + VectorParams, + Distance, + PointStruct, +) +from pipeline.config import config + + +class AsyncQD: + def __init__(self): + self.client = AsyncQdrantClient( + host=config["qdrant_host"], + port=config["qdrant_port"], + api_key=config["qdrant_api_key"], + ) + + async def ensure_collection(self, name: str, dim: int): + """ + 如果 collection 不存在则创建。 + 适用于多知识库:每个知识库一个 collection。 + """ + try: + await self.client.get_collection(name) + except Exception: + await self.client.recreate_collection( + collection_name=name, + vectors_config=VectorParams(size=dim, distance=Distance.COSINE), + ) + + async def upsert_vectors(self, collection: str, vectors: list, payloads: list): + """ + 写入向量(id 自增或自行维护) + - vectors: [[float...], ...] + - payloads: [{"text": "...", "kb": "..."}] + """ + points = [PointStruct(id=i, vector=vectors[i], payload=payloads[i]) for i in range(len(vectors))] + + await self.client.upsert(collection_name=collection, points=points) + + async def vector_search(self, collection: str, query_vector: list, top_k: int = 5): + """ + 查询向量相似度 + """ + result = await self.client.search( + collection_name=collection, + query_vector=query_vector, + limit=top_k, + ) + return result + + +qd_client = AsyncQD() diff --git a/src/pipeline/core/qdrant_client.py b/src/pipeline/core/qdrant_client.py deleted file mode 100644 index e69de29..0000000