已开启
为DocsGPT 新增openGauss DataVec作为向量数据库选项 #101
HLY-cloud创建于 5月31日
为DocsGPT 新增openGauss DataVec作为向量数据库选项 #101
已开启
共 7 个文件变更+1059-0
| @@ -0,0 +1,56 @@ | |||
| 1 | +## 项目说明 | ||
| 2 | + | ||
| 3 | +该项目是为[DocsGPT](https://github.com/arc53/DocsGPT) 提供openGauss DataVec作为向量数据库的选项。 | ||
| 4 | + | ||
| 5 | +**DocsGPT** 是一个开源的本地私有知识库问答(Q&A)系统,在 GitHub 上已获得超过 17k Star。 | ||
| 6 | + | ||
| 7 | +它支持用户将私有文档(如 PDF、DOCX、TXT、网页、本地代码库等)进行本地化索引,并通过自然语言进行智能提问。系统会自动检索相关文本段落,并结合大语言模型(LLM)给出精准、带引用源的回答。 | ||
| 8 | + | ||
| 9 | +目前,DocsGPT已经支持pgvector, milvus等主流向量数据库,但还不能支持openGauss DataVec,所以有了该项目。 | ||
| 10 | + | ||
| 11 | + | ||
| 12 | + | ||
| 13 | +## 项目结构 | ||
| 14 | + | ||
| 15 | +``` | ||
| 16 | +DocsGPT/ | ||
| 17 | +├── opengauss_datavec.py # 核心实现:openGauss DataVec 向量存储适配器 | ||
| 18 | +├── vector_creator.py # 修改:注册 opengauss_datavec 存储类型 | ||
| 19 | +├── settings.py # 修改:新增 OPENGAUSS_CONNECTION_STRING 配置项 | ||
| 20 | +├── test_opengauss_datavec.py # 单元测试 | ||
| 21 | +├── 实践教程.md # 部署与使用教程 | ||
| 22 | +└── README.md # 本文件 | ||
| 23 | +``` | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +## 设计说明 | ||
| 28 | + | ||
| 29 | +所有文档共用一张表 `documents`。由于不同文档的元数据字段各异,采用 JSONB 类型统一存储: | ||
| 30 | + | ||
| 31 | +```sql | ||
| 32 | +CREATE TABLE IF NOT EXISTS documents ( | ||
| 33 | + id BIGSERIAL PRIMARY KEY, | ||
| 34 | + text TEXT NOT NULL, | ||
| 35 | + embedding vector(xxx), -- 维度由 embedding 模型动态决定 | ||
| 36 | + metadata JSONB, | ||
| 37 | + source_id TEXT NOT NULL, | ||
| 38 | + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP | ||
| 39 | +); | ||
| 40 | +``` | ||
| 41 | + | ||
| 42 | +### 索引策略 | ||
| 43 | + | ||
| 44 | +建立 3 个索引,分别服务于向量检索和等值过滤: | ||
| 45 | + | ||
| 46 | +```sql | ||
| 47 | +-- 向量近邻检索(IVFFlat,lists=100) | ||
| 48 | +CREATE INDEX documents_embedding_ivfflat_idx | ||
| 49 | + ON documents USING ivfflat (embedding vector_l2_ops) WITH (lists = 100); | ||
| 50 | + | ||
| 51 | +-- 按文档源过滤(删除、查询时按 source_id 定位) | ||
| 52 | +CREATE INDEX documents_source_id_idx ON documents (source_id); | ||
| 53 | + | ||
| 54 | +-- 主键 id 自带 B-tree 索引,用于单条 chunk 删除 | ||
| 55 | +``` | ||
| 56 | + | ||
| @@ -0,0 +1,285 @@ | |||
| 1 | +import json | ||
| 2 | +import logging | ||
| 3 | +import math | ||
| 4 | +from typing import Any, Dict, List, Optional | ||
| 5 | + | ||
| 6 | +from application.core.settings import settings | ||
| 7 | +from application.vectorstore.base import BaseVectorStore | ||
| 8 | +from application.vectorstore.document_class import Document | ||
| 9 | + | ||
| 10 | +_TABLE = "documents" | ||
| 11 | +_INSERT_BATCH_SIZE = 500 | ||
| 12 | +_CONNECT_KWARGS = dict( | ||
| 13 | + keepalives=1, | ||
| 14 | + keepalives_idle=30, | ||
| 15 | + keepalives_interval=10, | ||
| 16 | + keepalives_count=5, | ||
| 17 | +) | ||
| 18 | + | ||
| 19 | + | ||
| 20 | +class OpenGaussDataVecStore(BaseVectorStore): | ||
| 21 | + """Vector store backed by openGauss DataVec. | ||
| 22 | + | ||
| 23 | + Single shared table 'documents' with fixed schema: | ||
| 24 | + id, text, embedding (vector), metadata (jsonb), source_id, created_at | ||
| 25 | + | ||
| 26 | + Requires OPENGAUSS_CONNECTION_STRING in settings. | ||
| 27 | + """ | ||
| 28 | + | ||
| 29 | + # Set once on first instantiation via _load_driver() | ||
| 30 | + _psycopg2 = None | ||
| 31 | + _sql = None | ||
| 32 | + _pg_extras = None | ||
| 33 | + | ||
| 34 | + | ||
| 35 | + def _load_driver(cls): | ||
| 36 | + if cls._psycopg2 is not None: | ||
| 37 | + return | ||
| 38 | + try: | ||
| 39 | + import psycopg2 | ||
| 40 | + import psycopg2.sql | ||
| 41 | + from psycopg2 import extras | ||
| 42 | + except ImportError: | ||
| 43 | + raise ImportError( | ||
| 44 | + "psycopg2 is required for openGauss. " | ||
| 45 | + "Install with: pip install psycopg2-binary" | ||
| 46 | + ) | ||
| 47 | + cls._psycopg2 = psycopg2 | ||
| 48 | + cls._sql = psycopg2.sql | ||
| 49 | + cls._pg_extras = extras | ||
| 50 | + | ||
| 51 | + def __init__(self, source_id: str = "", embeddings_key: str = "embeddings"): | ||
| 52 | + super().__init__() | ||
| 53 | + self._load_driver() | ||
| 54 | + self._source_id = str(source_id) | ||
| 55 | + self._embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key) | ||
| 56 | + self._embedding_dimension = self._resolve_embedding_dimension() | ||
| 57 | + self._conn_str = getattr(settings, "OPENGAUSS_CONNECTION_STRING", None) | ||
| 58 | + if not self._conn_str: | ||
| 59 | + raise ValueError("OPENGAUSS_CONNECTION_STRING is required in settings.") | ||
| 60 | + self._connection = None | ||
| 61 | + self._ensure_table_exists() | ||
| 62 | + | ||
| 63 | + # ------------------------------------------------------------------ | ||
| 64 | + # Internal helpers | ||
| 65 | + # ------------------------------------------------------------------ | ||
| 66 | + | ||
| 67 | + def _get_connection(self): | ||
| 68 | + """Return the shared connection, reconnecting if closed or stale.""" | ||
| 69 | + if self._connection is not None and not self._connection.closed: | ||
| 70 | + try: | ||
| 71 | + self._connection.poll() | ||
| 72 | + return self._connection | ||
| 73 | + except Exception: | ||
| 74 | + pass | ||
| 75 | + self._connection = self._psycopg2.connect(self._conn_str, **_CONNECT_KWARGS) | ||
| 76 | + return self._connection | ||
| 77 | + | ||
| 78 | + | ||
| 79 | + def _vec_to_str(vec) -> str: | ||
| 80 | + floats = [float(v) for v in vec] | ||
| 81 | + if any(math.isnan(f) or math.isinf(f) for f in floats): | ||
| 82 | + raise ValueError("Vector contains NaN or Inf values") | ||
| 83 | + return "[" + ",".join(str(f) for f in floats) + "]" | ||
| 84 | + | ||
| 85 | + | ||
| 86 | + def _parse_metadata(raw) -> dict: | ||
| 87 | + if isinstance(raw, dict): | ||
| 88 | + return raw | ||
| 89 | + if raw: | ||
| 90 | + try: | ||
| 91 | + return json.loads(raw) | ||
| 92 | + except (TypeError, ValueError): | ||
| 93 | + pass | ||
| 94 | + return {} | ||
| 95 | + | ||
| 96 | + def _resolve_embedding_dimension(self) -> int: | ||
| 97 | + """Resolve the actual embedding dimension via a probe call.""" | ||
| 98 | + probe = self._embedding.embed_query("dimension probe") | ||
| 99 | + actual_dim = len(probe) | ||
| 100 | + if actual_dim <= 0: | ||
| 101 | + raise ValueError("Embedding probe returned an empty vector") | ||
| 102 | + | ||
| 103 | + declared = getattr(self._embedding, "dimension", None) | ||
| 104 | + if declared != actual_dim: | ||
| 105 | + logging.warning( | ||
| 106 | + "Embedding dimension mismatch: declared=%s actual=%s. Using actual.", | ||
| 107 | + declared, | ||
| 108 | + actual_dim, | ||
| 109 | + ) | ||
| 110 | + try: | ||
| 111 | + self._embedding.dimension = actual_dim | ||
| 112 | + except Exception: | ||
| 113 | + pass | ||
| 114 | + return actual_dim | ||
| 115 | + | ||
| 116 | + def _ensure_table_exists(self): | ||
| 117 | + sql = self._sql | ||
| 118 | + table = sql.Identifier(_TABLE) | ||
| 119 | + conn = self._get_connection() | ||
| 120 | + with conn, conn.cursor() as cur: | ||
| 121 | + cur.execute( | ||
| 122 | + sql.SQL(""" | ||
| 123 | + CREATE TABLE IF NOT EXISTS {table} ( | ||
| 124 | + id BIGSERIAL PRIMARY KEY, | ||
| 125 | + text TEXT NOT NULL, | ||
| 126 | + embedding vector({dim}), | ||
| 127 | + metadata JSONB, | ||
| 128 | + source_id TEXT NOT NULL, | ||
| 129 | + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP | ||
| 130 | + ); | ||
| 131 | + """).format(table=table, dim=sql.Literal(int(self._embedding_dimension))) | ||
| 132 | + ) | ||
| 133 | + cur.execute( | ||
| 134 | + sql.SQL(""" | ||
| 135 | + CREATE INDEX IF NOT EXISTS {idx} | ||
| 136 | + ON {table} USING ivfflat (embedding vector_l2_ops) WITH (lists = 100); | ||
| 137 | + """).format( | ||
| 138 | + idx=sql.Identifier(f"{_TABLE}_embedding_ivfflat_idx"), | ||
| 139 | + table=table, | ||
| 140 | + ) | ||
| 141 | + ) | ||
| 142 | + cur.execute( | ||
| 143 | + sql.SQL(""" | ||
| 144 | + CREATE INDEX IF NOT EXISTS {idx} | ||
| 145 | + ON {table} (source_id); | ||
| 146 | + """).format( | ||
| 147 | + idx=sql.Identifier(f"{_TABLE}_source_id_idx"), | ||
| 148 | + table=table, | ||
| 149 | + ) | ||
| 150 | + ) | ||
| 151 | + | ||
| 152 | + # ------------------------------------------------------------------ | ||
| 153 | + # BaseVectorStore interface | ||
| 154 | + # ------------------------------------------------------------------ | ||
| 155 | + | ||
| 156 | + def search(self, question: str, k: int = 2, *args, **kwargs) -> List[Document]: | ||
| 157 | + sql = self._sql | ||
| 158 | + query_vec = self._embedding.embed_query(question) | ||
| 159 | + try: | ||
| 160 | + conn = self._get_connection() | ||
| 161 | + with conn, conn.cursor() as cur: | ||
| 162 | + cur.execute( | ||
| 163 | + sql.SQL(""" | ||
| 164 | + SELECT text, metadata | ||
| 165 | + FROM {table} | ||
| 166 | + WHERE source_id = %s | ||
| 167 | + ORDER BY embedding <-> %s::vector | ||
| 168 | + LIMIT %s; | ||
| 169 | + """).format(table=sql.Identifier(_TABLE)), | ||
| 170 | + (self._source_id, self._vec_to_str(query_vec), k), | ||
| 171 | + ) | ||
| 172 | + return [ | ||
| 173 | + Document( | ||
| 174 | + page_content=text, metadata=self._parse_metadata(meta) | ||
| 175 | + ) | ||
| 176 | + for text, meta in cur.fetchall() | ||
| 177 | + ] | ||
| 178 | + except Exception as e: | ||
| 179 | + logging.error(f"OpenGaussDataVecStore.search error: {e}", exc_info=True) | ||
| 180 | + return [] | ||
| 181 | + | ||
| 182 | + def add_texts( | ||
| 183 | + self, | ||
| 184 | + texts: List[str], | ||
| 185 | + metadatas: Optional[List[Dict[str, Any]]] = None, | ||
| 186 | + *args, | ||
| 187 | + **kwargs, | ||
| 188 | + ) -> List[str]: | ||
| 189 | + if not texts: | ||
| 190 | + return [] | ||
| 191 | + sql = self._sql | ||
| 192 | + embeddings = self._embedding.embed_documents(texts) | ||
| 193 | + metadatas = metadatas or [{}] * len(texts) | ||
| 194 | + rows = [ | ||
| 195 | + (text, self._vec_to_str(emb), json.dumps(meta), self._source_id) | ||
| 196 | + for text, emb, meta in zip(texts, embeddings, metadatas) | ||
| 197 | + ] | ||
| 198 | + | ||
| 199 | + query = sql.SQL(""" | ||
| 200 | + INSERT INTO {table} (text, embedding, metadata, source_id) | ||
| 201 | + VALUES %s | ||
| 202 | + RETURNING id; | ||
| 203 | + """).format(table=sql.Identifier(_TABLE)) | ||
| 204 | + | ||
| 205 | + conn = self._get_connection() | ||
| 206 | + inserted_ids = [] | ||
| 207 | + with conn, conn.cursor() as cur: | ||
| 208 | + for i in range(0, len(rows), _INSERT_BATCH_SIZE): | ||
| 209 | + self._pg_extras.execute_values( | ||
| 210 | + cur, | ||
| 211 | + query.as_string(cur), | ||
| 212 | + rows[i : i + _INSERT_BATCH_SIZE], | ||
| 213 | + template="(%s, %s::vector, %s::jsonb, %s)", | ||
| 214 | + fetch=True, | ||
| 215 | + ) | ||
| 216 | + inserted_ids.extend(str(row[0]) for row in cur.fetchall()) | ||
| 217 | + return inserted_ids | ||
| 218 | + | ||
| 219 | + def delete_index(self, *args, **kwargs): | ||
| 220 | + sql = self._sql | ||
| 221 | + conn = self._get_connection() | ||
| 222 | + with conn, conn.cursor() as cur: | ||
| 223 | + cur.execute( | ||
| 224 | + sql.SQL("DELETE FROM {table} WHERE source_id = %s;").format( | ||
| 225 | + table=sql.Identifier(_TABLE) | ||
| 226 | + ), | ||
| 227 | + (self._source_id,), | ||
| 228 | + ) | ||
| 229 | + | ||
| 230 | + def save_local(self, *args, **kwargs): | ||
| 231 | + pass | ||
| 232 | + | ||
| 233 | + def get_chunks(self) -> List[Dict[str, Any]]: | ||
| 234 | + sql = self._sql | ||
| 235 | + try: | ||
| 236 | + conn = self._get_connection() | ||
| 237 | + with conn, conn.cursor() as cur: | ||
| 238 | + cur.execute( | ||
| 239 | + sql.SQL( | ||
| 240 | + "SELECT id, text, metadata FROM {table} WHERE source_id = %s;" | ||
| 241 | + ).format(table=sql.Identifier(_TABLE)), | ||
| 242 | + (self._source_id,), | ||
| 243 | + ) | ||
| 244 | + return [ | ||
| 245 | + { | ||
| 246 | + "doc_id": str(doc_id), | ||
| 247 | + "text": text, | ||
| 248 | + "metadata": self._parse_metadata(meta), | ||
| 249 | + } | ||
| 250 | + for doc_id, text, meta in cur.fetchall() | ||
| 251 | + ] | ||
| 252 | + except Exception as e: | ||
| 253 | + logging.error(f"OpenGaussDataVecStore.get_chunks error: {e}") | ||
| 254 | + return [] | ||
| 255 | + | ||
| 256 | + def add_chunk(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> str: | ||
| 257 | + return self.add_texts([text], [metadata or {}])[0] | ||
| 258 | + | ||
| 259 | + def delete_chunk(self, chunk_id: str) -> bool: | ||
| 260 | + sql = self._sql | ||
| 261 | + try: | ||
| 262 | + conn = self._get_connection() | ||
| 263 | + with conn, conn.cursor() as cur: | ||
| 264 | + cur.execute( | ||
| 265 | + sql.SQL( | ||
| 266 | + "DELETE FROM {table} WHERE id = %s AND source_id = %s;" | ||
| 267 | + ).format(table=sql.Identifier(_TABLE)), | ||
| 268 | + (int(chunk_id), self._source_id), | ||
| 269 | + ) | ||
| 270 | + return cur.rowcount > 0 | ||
| 271 | + except Exception as e: | ||
| 272 | + logging.error(f"OpenGaussDataVecStore.delete_chunk error: {e}") | ||
| 273 | + return False | ||
| 274 | + | ||
| 275 | + def __enter__(self): | ||
| 276 | + return self | ||
| 277 | + | ||
| 278 | + def __exit__(self, *_): | ||
| 279 | + if self._connection and not self._connection.closed: | ||
| 280 | + self._connection.close() | ||
| 281 | + return False | ||
| 282 | + | ||
| 283 | + def __del__(self): | ||
| 284 | + if hasattr(self, "_connection") and self._connection and not self._connection.closed: | ||
| 285 | + self._connection.close() | ||
| @@ -0,0 +1,314 @@ | |||
| 1 | +import os | ||
| 2 | +from pathlib import Path | ||
| 3 | +from typing import Optional | ||
| 4 | + | ||
| 5 | +from pydantic import field_validator | ||
| 6 | +from pydantic_settings import BaseSettings, SettingsConfigDict | ||
| 7 | + | ||
| 8 | +current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +from application.core.db_uri import ( # noqa: E402 | ||
| 12 | + normalize_pgvector_connection_string, | ||
| 13 | + normalize_postgres_uri, | ||
| 14 | +) | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +class Settings(BaseSettings): | ||
| 18 | + model_config = SettingsConfigDict(extra="ignore") | ||
| 19 | + | ||
| 20 | + AUTH_TYPE: Optional[str] = None # simple_jwt, session_jwt, or None | ||
| 21 | + LLM_PROVIDER: str = "docsgpt" | ||
| 22 | + LLM_NAME: Optional[str] = None # if LLM_PROVIDER is openai, LLM_NAME can be gpt-4 or gpt-3.5-turbo | ||
| 23 | + EMBEDDINGS_NAME: str = "huggingface_sentence-transformers/all-mpnet-base-v2" | ||
| 24 | + EMBEDDINGS_BASE_URL: Optional[str] = None # Remote embeddings API URL (OpenAI-compatible) | ||
| 25 | + EMBEDDINGS_KEY: Optional[str] = None # api key for embeddings (if using openai, just copy API_KEY) | ||
| 26 | + # Optional directory of operator-supplied model YAMLs, loaded after the | ||
| 27 | + # built-in catalog under application/core/models/. Later wins on | ||
| 28 | + # duplicate model id. See application/core/models/README.md. | ||
| 29 | + MODELS_CONFIG_DIR: Optional[str] = None | ||
| 30 | + | ||
| 31 | + CELERY_BROKER_URL: str = "redis://localhost:6379/0" | ||
| 32 | + CELERY_RESULT_BACKEND: str = "redis://localhost:6379/1" | ||
| 33 | + # Prefetch=1 caps SIGKILL loss to one task. Visibility timeout must exceed | ||
| 34 | + # the longest legitimate task runtime (ingest, agent webhook) but stay | ||
| 35 | + # short enough that SIGKILLed tasks redeliver promptly. 1h matches Onyx | ||
| 36 | + # and Dify defaults; long ingests can override via env. | ||
| 37 | + CELERY_WORKER_PREFETCH_MULTIPLIER: int = 1 | ||
| 38 | + CELERY_VISIBILITY_TIMEOUT: int = 3600 | ||
| 39 | + # Recycle the prefork worker child once its resident size crosses this many | ||
| 40 | + # kilobytes — backstops native-heap growth from docling/torch parsing. 0 disables. | ||
| 41 | + CELERY_WORKER_MAX_MEMORY_PER_CHILD: int = 4194304 | ||
| 42 | + # Recycle the child after this many tasks; 0 disables (memory cap is the primary knob). | ||
| 43 | + CELERY_WORKER_MAX_TASKS_PER_CHILD: int = 0 | ||
| 44 | + # Only consulted when VECTOR_STORE=mongodb or when running scripts/db/backfill.py; user data lives in Postgres. | ||
| 45 | + MONGO_URI: Optional[str] = None | ||
| 46 | + # User-data Postgres DB. | ||
| 47 | + POSTGRES_URI: Optional[str] = None | ||
| 48 | + # On app startup, apply pending Alembic migrations. Default ON for dev; disable in prod if you manage schema out-of-band. | ||
| 49 | + AUTO_MIGRATE: bool = True | ||
| 50 | + # On app startup, create the target Postgres database if it's missing (requires CREATEDB privilege). Dev-friendly default. | ||
| 51 | + AUTO_CREATE_DB: bool = True | ||
| 52 | + LLM_PATH: str = os.path.join(current_dir, "models/docsgpt-7b-f16.gguf") | ||
| 53 | + DEFAULT_MAX_HISTORY: int = 150 | ||
| 54 | + DEFAULT_LLM_TOKEN_LIMIT: int = 128000 # Fallback when model not found in registry | ||
| 55 | + RESERVED_TOKENS: dict = { | ||
| 56 | + "system_prompt": 500, | ||
| 57 | + "current_query": 500, | ||
| 58 | + "safety_buffer": 1000, | ||
| 59 | + } | ||
| 60 | + DEFAULT_AGENT_LIMITS: dict = { | ||
| 61 | + "token_limit": 50000, | ||
| 62 | + "request_limit": 500, | ||
| 63 | + } | ||
| 64 | + UPLOAD_FOLDER: str = "inputs" | ||
| 65 | + PARSE_PDF_AS_IMAGE: bool = False | ||
| 66 | + PARSE_IMAGE_REMOTE: bool = False | ||
| 67 | + DOCLING_OCR_ENABLED: bool = False # Enable OCR for docling parsers (PDF, images) | ||
| 68 | + DOCLING_OCR_ATTACHMENTS_ENABLED: bool = False # Enable OCR for docling when parsing attachments | ||
| 69 | + # Pages docling's threaded pipeline buffers in flight; the library | ||
| 70 | + # default (100) drives worker RSS to ~3 GB on a mid-size PDF. | ||
| 71 | + DOCLING_PIPELINE_QUEUE_MAX_SIZE: int = 2 | ||
| 72 | + VECTOR_STORE: str = "faiss" # "faiss" or "elasticsearch" or "qdrant" or "milvus" or "lancedb" or "pgvector" or "opengauss_datavec" | ||
| 73 | + RETRIEVERS_ENABLED: list = ["classic_rag"] | ||
| 74 | + AGENT_NAME: str = "classic" | ||
| 75 | + FALLBACK_LLM_PROVIDER: Optional[str] = None # provider for fallback llm | ||
| 76 | + FALLBACK_LLM_NAME: Optional[str] = None # model name for fallback llm | ||
| 77 | + FALLBACK_LLM_API_KEY: Optional[str] = None # api key for fallback llm | ||
| 78 | + | ||
| 79 | + # Google Drive integration | ||
| 80 | + GOOGLE_CLIENT_ID: Optional[str] = None # Replace with your actual Google OAuth client ID | ||
| 81 | + GOOGLE_CLIENT_SECRET: Optional[str] = None # Replace with your actual Google OAuth client secret | ||
| 82 | + CONNECTOR_REDIRECT_BASE_URI: Optional[str] = ( | ||
| 83 | + "http://127.0.0.1:7091/api/connectors/callback" ##add redirect url as it is to your provider's console(gcp) | ||
| 84 | + ) | ||
| 85 | + | ||
| 86 | + # Microsoft Entra ID (Azure AD) integration | ||
| 87 | + MICROSOFT_CLIENT_ID: Optional[str] = None # Azure AD Application (client) ID | ||
| 88 | + MICROSOFT_CLIENT_SECRET: Optional[str] = None # Azure AD Application client secret | ||
| 89 | + MICROSOFT_TENANT_ID: Optional[str] = "common" # Azure AD Tenant ID (or 'common' for multi-tenant) | ||
| 90 | + MICROSOFT_AUTHORITY: Optional[str] = None # e.g., "https://login.microsoftonline.com/{tenant_id}" | ||
| 91 | + | ||
| 92 | + # Confluence Cloud integration | ||
| 93 | + CONFLUENCE_CLIENT_ID: Optional[str] = None | ||
| 94 | + CONFLUENCE_CLIENT_SECRET: Optional[str] = None | ||
| 95 | + | ||
| 96 | + # GitHub source | ||
| 97 | + GITHUB_ACCESS_TOKEN: Optional[str] = None # PAT token with read repo access | ||
| 98 | + | ||
| 99 | + # LLM Cache | ||
| 100 | + CACHE_REDIS_URL: str = "redis://localhost:6379/2" | ||
| 101 | + | ||
| 102 | + API_URL: str = "http://localhost:7091" # backend url for celery worker | ||
| 103 | + MCP_OAUTH_REDIRECT_URI: Optional[str] = None # public callback URL for MCP OAuth | ||
| 104 | + INTERNAL_KEY: Optional[str] = None # internal api key for worker-to-backend auth | ||
| 105 | + | ||
| 106 | + API_KEY: Optional[str] = None # LLM api key (used by LLM_PROVIDER) | ||
| 107 | + | ||
| 108 | + # Provider-specific API keys (for multi-model support) | ||
| 109 | + OPENAI_API_KEY: Optional[str] = None | ||
| 110 | + ANTHROPIC_API_KEY: Optional[str] = None | ||
| 111 | + GOOGLE_API_KEY: Optional[str] = None | ||
| 112 | + GROQ_API_KEY: Optional[str] = None | ||
| 113 | + HUGGINGFACE_API_KEY: Optional[str] = None | ||
| 114 | + OPEN_ROUTER_API_KEY: Optional[str] = None | ||
| 115 | + NOVITA_API_KEY: Optional[str] = None | ||
| 116 | + | ||
| 117 | + OPENAI_API_BASE: Optional[str] = None # azure openai api base url | ||
| 118 | + OPENAI_API_VERSION: Optional[str] = None # azure openai api version | ||
| 119 | + AZURE_DEPLOYMENT_NAME: Optional[str] = None # azure deployment name for answering | ||
| 120 | + AZURE_EMBEDDINGS_DEPLOYMENT_NAME: Optional[str] = None # azure deployment name for embeddings | ||
| 121 | + OPENAI_BASE_URL: Optional[str] = None # openai base url for open ai compatable models | ||
| 122 | + | ||
| 123 | + # elasticsearch | ||
| 124 | + ELASTIC_CLOUD_ID: Optional[str] = None # cloud id for elasticsearch | ||
| 125 | + ELASTIC_USERNAME: Optional[str] = None # username for elasticsearch | ||
| 126 | + ELASTIC_PASSWORD: Optional[str] = None # password for elasticsearch | ||
| 127 | + ELASTIC_URL: Optional[str] = None # url for elasticsearch | ||
| 128 | + ELASTIC_INDEX: Optional[str] = "docsgpt" # index name for elasticsearch | ||
| 129 | + | ||
| 130 | + # SageMaker config | ||
| 131 | + SAGEMAKER_ENDPOINT: Optional[str] = None # SageMaker endpoint name | ||
| 132 | + SAGEMAKER_REGION: Optional[str] = None # SageMaker region name | ||
| 133 | + SAGEMAKER_ACCESS_KEY: Optional[str] = None # SageMaker access key | ||
| 134 | + SAGEMAKER_SECRET_KEY: Optional[str] = None # SageMaker secret key | ||
| 135 | + | ||
| 136 | + # prem ai project id | ||
| 137 | + PREMAI_PROJECT_ID: Optional[str] = None | ||
| 138 | + | ||
| 139 | + # Qdrant vectorstore config | ||
| 140 | + QDRANT_COLLECTION_NAME: Optional[str] = "docsgpt" | ||
| 141 | + QDRANT_LOCATION: Optional[str] = None | ||
| 142 | + QDRANT_URL: Optional[str] = None | ||
| 143 | + QDRANT_PORT: Optional[int] = 6333 | ||
| 144 | + QDRANT_GRPC_PORT: int = 6334 | ||
| 145 | + QDRANT_PREFER_GRPC: bool = False | ||
| 146 | + QDRANT_HTTPS: Optional[bool] = None | ||
| 147 | + QDRANT_API_KEY: Optional[str] = None | ||
| 148 | + QDRANT_PREFIX: Optional[str] = None | ||
| 149 | + QDRANT_TIMEOUT: Optional[float] = None | ||
| 150 | + QDRANT_HOST: Optional[str] = None | ||
| 151 | + QDRANT_PATH: Optional[str] = None | ||
| 152 | + QDRANT_DISTANCE_FUNC: str = "Cosine" | ||
| 153 | + | ||
| 154 | + # PGVector vectorstore config. Write the URI in whichever form you | ||
| 155 | + # prefer — ``postgres://``, ``postgresql://``, or even the SQLAlchemy | ||
| 156 | + # dialect form (``postgresql+psycopg://``) are all accepted and | ||
| 157 | + # normalized internally for ``psycopg.connect()``. | ||
| 158 | + PGVECTOR_CONNECTION_STRING: Optional[str] = None | ||
| 159 | + | ||
| 160 | + # openGauss DataVec vectorstore config. | ||
| 161 | + # DataVec is built into openGauss — no extension installation needed. | ||
| 162 | + # Use a libpq-style DSN, e.g.: | ||
| 163 | + # "host=127.0.0.1 port=5432 dbname=mydb user=myuser password=xxx" | ||
| 164 | + OPENGAUSS_CONNECTION_STRING: Optional[str] = None | ||
| 165 | + | ||
| 166 | + # Milvus vectorstore config | ||
| 167 | + MILVUS_COLLECTION_NAME: Optional[str] = "docsgpt" | ||
| 168 | + MILVUS_URI: Optional[str] = "./milvus_local.db" # milvus lite version as default | ||
| 169 | + MILVUS_TOKEN: Optional[str] = "" | ||
| 170 | + | ||
| 171 | + # LanceDB vectorstore config | ||
| 172 | + LANCEDB_PATH: str = "./data/lancedb" # Path where LanceDB stores its local data | ||
| 173 | + LANCEDB_TABLE_NAME: Optional[str] = "docsgpts" # Name of the table to use for storing vectors | ||
| 174 | + | ||
| 175 | + FLASK_DEBUG_MODE: bool = False | ||
| 176 | + STORAGE_TYPE: str = "local" # local or s3 | ||
| 177 | + | ||
| 178 | + # Anonymous startup version check for security issues. | ||
| 179 | + VERSION_CHECK: bool = True | ||
| 180 | + URL_STRATEGY: str = "backend" # backend or s3 | ||
| 181 | + | ||
| 182 | + JWT_SECRET_KEY: str = "" | ||
| 183 | + | ||
| 184 | + # Encryption settings | ||
| 185 | + ENCRYPTION_SECRET_KEY: str = "default-docsgpt-encryption-key" | ||
| 186 | + | ||
| 187 | + TTS_PROVIDER: str = "google_tts" # google_tts or elevenlabs | ||
| 188 | + ELEVENLABS_API_KEY: Optional[str] = None | ||
| 189 | + STT_PROVIDER: str = "openai" # openai or faster_whisper | ||
| 190 | + OPENAI_STT_MODEL: str = "gpt-4o-mini-transcribe" | ||
| 191 | + STT_LANGUAGE: Optional[str] = None | ||
| 192 | + STT_MAX_FILE_SIZE_MB: int = 50 | ||
| 193 | + STT_ENABLE_TIMESTAMPS: bool = False | ||
| 194 | + STT_ENABLE_DIARIZATION: bool = False | ||
| 195 | + | ||
| 196 | + # Tool pre-fetch settings | ||
| 197 | + ENABLE_TOOL_PREFETCH: bool = True | ||
| 198 | + | ||
| 199 | + # Config-free tools on by default in agentless chats. ``scheduler`` is | ||
| 200 | + # dual-registered (also in ``BUILTIN_AGENT_TOOLS``) so the same synthetic id | ||
| 201 | + # resolves whether reached via defaults or the agent picker. | ||
| 202 | + DEFAULT_CHAT_TOOLS: list = ["memory", "read_webpage", "scheduler"] | ||
| 203 | + | ||
| 204 | + # Conversation Compression Settings | ||
| 205 | + ENABLE_CONVERSATION_COMPRESSION: bool = True | ||
| 206 | + COMPRESSION_THRESHOLD_PERCENTAGE: float = 0.8 # Trigger at 80% of context | ||
| 207 | + COMPRESSION_MODEL_OVERRIDE: Optional[str] = None # Use different model for compression | ||
| 208 | + COMPRESSION_PROMPT_VERSION: str = "v1.0" # Track prompt iterations | ||
| 209 | + COMPRESSION_MAX_HISTORY_POINTS: int = 3 # Keep only last N compression points to prevent DB bloat | ||
| 210 | + | ||
| 211 | + # Internal SSE push channel (notifications + durable replay journal) | ||
| 212 | + # Master switch — when False, /api/events emits a "push_disabled" comment | ||
| 213 | + # and returns; clients fall back to polling. Publisher becomes a no-op. | ||
| 214 | + ENABLE_SSE_PUSH: bool = True | ||
| 215 | + # Per-user durable backlog cap (~entries). At typical event rates this | ||
| 216 | + # gives ~24h of replay; tune up for verbose feeds, down for memory. | ||
| 217 | + EVENTS_STREAM_MAXLEN: int = 1000 | ||
| 218 | + # SSE keepalive comment cadence. Must sit under Cloudflare's 100s idle | ||
| 219 | + # close and iOS Safari's ~60s — 15s gives generous headroom. | ||
| 220 | + SSE_KEEPALIVE_SECONDS: int = 15 | ||
| 221 | + # Cap on simultaneous SSE connections per user. Each connection holds | ||
| 222 | + # one WSGI thread (32 per gunicorn worker) and one Redis pub/sub | ||
| 223 | + # connection. 8 covers normal multi-tab use without letting one user | ||
| 224 | + # starve the pool. Set to 0 to disable the cap. | ||
| 225 | + SSE_MAX_CONCURRENT_PER_USER: int = 8 | ||
| 226 | + # Per-request cap on the number of backlog entries XRANGE returns | ||
| 227 | + # for ``/api/events`` snapshots. Bounds the bytes a single replay | ||
| 228 | + # can move from Redis to the wire — a malicious client looping | ||
| 229 | + # ``Last-Event-ID=<oldest>`` reconnects can only enumerate this | ||
| 230 | + # many entries per round-trip. Combined with the per-user | ||
| 231 | + # connection cap above and the windowed budget below, total | ||
| 232 | + # enumeration throughput is bounded. | ||
| 233 | + EVENTS_REPLAY_MAX_PER_REQUEST: int = 200 | ||
| 234 | + # Sliding-window cap on snapshot replays per user. Once the budget | ||
| 235 | + # is exhausted the route returns HTTP 429 with the cursor pinned; | ||
| 236 | + # the client backs off and retries after the window rolls over. | ||
| 237 | + EVENTS_REPLAY_BUDGET_REQUESTS_PER_WINDOW: int = 30 | ||
| 238 | + EVENTS_REPLAY_BUDGET_WINDOW_SECONDS: int = 60 | ||
| 239 | + | ||
| 240 | + # Retention for the ``message_events`` journal. The ``cleanup_message_events`` | ||
| 241 | + # beat task deletes rows older than this. Reconnect-replay only | ||
| 242 | + # needs the journal for streams a client could still be tailing, | ||
| 243 | + # so 14 days is a generous default that covers paused/tool-action | ||
| 244 | + # flows without unbounded table growth. | ||
| 245 | + MESSAGE_EVENTS_RETENTION_DAYS: int = 14 | ||
| 246 | + | ||
| 247 | + # Remote Device feature. | ||
| 248 | + REMOTE_DEVICE_SESSION_IDLE_SECONDS: int = 60 | ||
| 249 | + REMOTE_DEVICE_REQUIRE_SIGNATURE: bool = False | ||
| 250 | + REMOTE_DEVICE_PAIRING_TTL_SECONDS: int = 600 | ||
| 251 | + # Redis-backed broker tunables (route invocations cross-process so a | ||
| 252 | + # scheduled/Celery run reaches the web-held device session). The command | ||
| 253 | + # queue TTL must exceed the max command drain deadline (the tool caps | ||
| 254 | + # timeout_ms at 600s, drained with a +5s margin = 605s) so a queued command | ||
| 255 | + # for a briefly-offline device isn't evicted before its own drain gives up. | ||
| 256 | + REMOTE_DEVICE_CMD_QUEUE_TTL_SECONDS: int = 900 | ||
| 257 | + REMOTE_DEVICE_INVOCATION_TTL_SECONDS: int = 900 | ||
| 258 | + REMOTE_DEVICE_OUTPUT_STREAM_MAXLEN: int = 10_000 | ||
| 259 | + | ||
| 260 | + # Scheduler (see scheduler.md). | ||
| 261 | + SCHEDULE_DISPATCHER_INTERVAL: int = 30 | ||
| 262 | + SCHEDULE_MIN_INTERVAL: int = 900 | ||
| 263 | + SCHEDULE_MAX_PER_USER: int = 50 | ||
| 264 | + SCHEDULE_RUN_TIMEOUT: int = 600 | ||
| 265 | + SCHEDULE_MISFIRE_GRACE: int = 60 | ||
| 266 | + SCHEDULE_AUTOPAUSE_FAILURES: int = 3 | ||
| 267 | + SCHEDULE_ONCE_MAX_HORIZON: int = 31_536_000 | ||
| 268 | + SCHEDULE_RUN_OUTPUT_RETENTION_DAYS: int = 90 | ||
| 269 | + | ||
| 270 | + | ||
| 271 | + | ||
| 272 | + def _normalize_postgres_uri_validator(cls, v): | ||
| 273 | + return normalize_postgres_uri(v) | ||
| 274 | + | ||
| 275 | + | ||
| 276 | + | ||
| 277 | + def _normalize_pgvector_connection_string_validator(cls, v): | ||
| 278 | + return normalize_pgvector_connection_string(v) | ||
| 279 | + | ||
| 280 | + | ||
| 281 | + "API_KEY", | ||
| 282 | + "OPENAI_API_KEY", | ||
| 283 | + "ANTHROPIC_API_KEY", | ||
| 284 | + "GOOGLE_API_KEY", | ||
| 285 | + "GROQ_API_KEY", | ||
| 286 | + "HUGGINGFACE_API_KEY", | ||
| 287 | + "NOVITA_API_KEY", | ||
| 288 | + "EMBEDDINGS_KEY", | ||
| 289 | + "FALLBACK_LLM_API_KEY", | ||
| 290 | + "QDRANT_API_KEY", | ||
| 291 | + "ELEVENLABS_API_KEY", | ||
| 292 | + "INTERNAL_KEY", | ||
| 293 | + mode="before", | ||
| 294 | + ) | ||
| 295 | + | ||
| 296 | + def normalize_api_key(cls, v: Optional[str]) -> Optional[str]: | ||
| 297 | + """ | ||
| 298 | + Normalize API keys: convert 'None', 'none', empty strings, | ||
| 299 | + and whitespace-only strings to actual None. | ||
| 300 | + Handles Pydantic loading 'None' from .env as string "None". | ||
| 301 | + """ | ||
| 302 | + if v is None: | ||
| 303 | + return None | ||
| 304 | + if not isinstance(v, str): | ||
| 305 | + return v | ||
| 306 | + stripped = v.strip() | ||
| 307 | + if stripped == "" or stripped.lower() == "none": | ||
| 308 | + return None | ||
| 309 | + return stripped | ||
| 310 | + | ||
| 311 | + | ||
| 312 | +# Project root is one level above application/ | ||
| 313 | +path = Path(__file__).parent.parent.parent.absolute() | ||
| 314 | +settings = Settings(_env_file=path.joinpath(".env"), _env_file_encoding="utf-8") | ||
| @@ -0,0 +1,314 @@ | |||
| 1 | +from unittest.mock import MagicMock, Mock, patch | ||
| 2 | + | ||
| 3 | +import pytest | ||
| 4 | + | ||
| 5 | +from application.vectorstore.opengauss_datavec import OpenGaussDataVecStore | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +def _make_store(source_id="test-source", embeddings_key="key"): | ||
| 9 | + """Create an OpenGaussDataVecStore with all external deps mocked.""" | ||
| 10 | + mock_emb = Mock() | ||
| 11 | + mock_emb.embed_query = Mock(return_value=[0.1, 0.2, 0.3]) | ||
| 12 | + mock_emb.embed_documents = Mock(return_value=[[0.1, 0.2, 0.3]]) | ||
| 13 | + mock_emb.dimension = 768 | ||
| 14 | + | ||
| 15 | + with patch.object(OpenGaussDataVecStore, "_load_driver"), \ | ||
| 16 | + patch( | ||
| 17 | + "application.vectorstore.base.BaseVectorStore._get_embeddings", | ||
| 18 | + return_value=mock_emb, | ||
| 19 | + ), \ | ||
| 20 | + patch( | ||
| 21 | + "application.vectorstore.opengauss_datavec.settings" | ||
| 22 | + ) as mock_settings, \ | ||
| 23 | + patch.object(OpenGaussDataVecStore, "_ensure_table_exists"): | ||
| 24 | + | ||
| 25 | + mock_settings.EMBEDDINGS_NAME = "test_model" | ||
| 26 | + mock_settings.OPENGAUSS_CONNECTION_STRING = "host=localhost dbname=test" | ||
| 27 | + | ||
| 28 | + store = OpenGaussDataVecStore( | ||
| 29 | + source_id=source_id, embeddings_key=embeddings_key | ||
| 30 | + ) | ||
| 31 | + | ||
| 32 | + # Set driver mocks as instance attrs (won't leak to other tests) | ||
| 33 | + store._sql = MagicMock() | ||
| 34 | + store._pg_extras = MagicMock() | ||
| 35 | + | ||
| 36 | + # Wire up mock connection | ||
| 37 | + mock_conn = MagicMock() | ||
| 38 | + mock_cursor = MagicMock() | ||
| 39 | + mock_conn.cursor.return_value.__enter__ = Mock(return_value=mock_cursor) | ||
| 40 | + mock_conn.cursor.return_value.__exit__ = Mock(return_value=False) | ||
| 41 | + store._get_connection = Mock(return_value=mock_conn) | ||
| 42 | + | ||
| 43 | + return store, mock_conn, mock_cursor, mock_emb | ||
| 44 | + | ||
| 45 | + | ||
| 46 | + | ||
| 47 | +class TestOpenGaussDataVecStoreInit: | ||
| 48 | + def test_source_id_stored_as_is(self): | ||
| 49 | + store, _, _, _ = _make_store(source_id="abc123") | ||
| 50 | + assert store._source_id == "abc123" | ||
| 51 | + | ||
| 52 | + def test_missing_connection_string_raises(self): | ||
| 53 | + mock_emb = Mock(dimension=768, embed_query=Mock(return_value=[0.1, 0.2, 0.3])) | ||
| 54 | + | ||
| 55 | + with patch.object(OpenGaussDataVecStore, "_load_driver"), \ | ||
| 56 | + patch( | ||
| 57 | + "application.vectorstore.base.BaseVectorStore._get_embeddings", | ||
| 58 | + return_value=mock_emb, | ||
| 59 | + ), \ | ||
| 60 | + patch( | ||
| 61 | + "application.vectorstore.opengauss_datavec.settings" | ||
| 62 | + ) as mock_settings: | ||
| 63 | + | ||
| 64 | + mock_settings.EMBEDDINGS_NAME = "test_model" | ||
| 65 | + mock_settings.OPENGAUSS_CONNECTION_STRING = None | ||
| 66 | + | ||
| 67 | + with pytest.raises(ValueError, match="OPENGAUSS_CONNECTION_STRING"): | ||
| 68 | + OpenGaussDataVecStore(source_id="test", embeddings_key="key") | ||
| 69 | + | ||
| 70 | + | ||
| 71 | + | ||
| 72 | +class TestOpenGaussDataVecStoreSearch: | ||
| 73 | + def test_search_returns_documents(self): | ||
| 74 | + store, _, mock_cursor, mock_emb = _make_store() | ||
| 75 | + mock_cursor.fetchall.return_value = [ | ||
| 76 | + ("hello world", {"source": "test.txt"}), | ||
| 77 | + ("foo bar", {"source": "test2.txt"}), | ||
| 78 | + ] | ||
| 79 | + | ||
| 80 | + results = store.search("query", k=2) | ||
| 81 | + | ||
| 82 | + mock_emb.embed_query.assert_called_with("query") | ||
| 83 | + assert len(results) == 2 | ||
| 84 | + assert results[0].page_content == "hello world" | ||
| 85 | + assert results[0].metadata == {"source": "test.txt"} | ||
| 86 | + | ||
| 87 | + def test_search_returns_empty_on_error(self): | ||
| 88 | + store, _, mock_cursor, _ = _make_store() | ||
| 89 | + mock_cursor.execute.side_effect = Exception("connection lost") | ||
| 90 | + | ||
| 91 | + results = store.search("query") | ||
| 92 | + assert results == [] | ||
| 93 | + | ||
| 94 | + def test_search_handles_null_metadata(self): | ||
| 95 | + store, _, mock_cursor, _ = _make_store() | ||
| 96 | + mock_cursor.fetchall.return_value = [("text", None)] | ||
| 97 | + | ||
| 98 | + results = store.search("query") | ||
| 99 | + assert len(results) == 1 | ||
| 100 | + assert results[0].metadata == {} | ||
| 101 | + | ||
| 102 | + def test_search_handles_string_metadata(self): | ||
| 103 | + store, _, mock_cursor, _ = _make_store() | ||
| 104 | + mock_cursor.fetchall.return_value = [("text", '{"key": "val"}')] | ||
| 105 | + | ||
| 106 | + results = store.search("query") | ||
| 107 | + assert results[0].metadata == {"key": "val"} | ||
| 108 | + | ||
| 109 | + def test_search_filters_by_source_id(self): | ||
| 110 | + store, _, mock_cursor, _ = _make_store(source_id="src42") | ||
| 111 | + mock_cursor.fetchall.return_value = [] | ||
| 112 | + | ||
| 113 | + store.search("query") | ||
| 114 | + | ||
| 115 | + params = mock_cursor.execute.call_args[0][1] | ||
| 116 | + assert params[0] == "src42" | ||
| 117 | + | ||
| 118 | + | ||
| 119 | + | ||
| 120 | +class TestOpenGaussDataVecStoreAddTexts: | ||
| 121 | + def test_add_texts_inserts_and_returns_ids(self): | ||
| 122 | + store, _, mock_cursor, mock_emb = _make_store() | ||
| 123 | + mock_emb.embed_documents.return_value = [[0.1, 0.2], [0.3, 0.4]] | ||
| 124 | + mock_cursor.fetchall.return_value = [(1,), (2,)] | ||
| 125 | + | ||
| 126 | + ids = store.add_texts(["text1", "text2"], [{"a": 1}, {"b": 2}]) | ||
| 127 | + | ||
| 128 | + assert ids == ["1", "2"] | ||
| 129 | + store._pg_extras.execute_values.assert_called_once() | ||
| 130 | + | ||
| 131 | + def test_add_texts_empty_returns_empty(self): | ||
| 132 | + store, _, _, _ = _make_store() | ||
| 133 | + assert store.add_texts([]) == [] | ||
| 134 | + | ||
| 135 | + def test_add_texts_default_metadatas(self): | ||
| 136 | + store, _, mock_cursor, mock_emb = _make_store() | ||
| 137 | + mock_emb.embed_documents.return_value = [[0.1, 0.2]] | ||
| 138 | + mock_cursor.fetchall.return_value = [(1,)] | ||
| 139 | + | ||
| 140 | + ids = store.add_texts(["text1"]) | ||
| 141 | + assert ids == ["1"] | ||
| 142 | + | ||
| 143 | + def test_add_texts_passes_source_id(self): | ||
| 144 | + store, _, mock_cursor, mock_emb = _make_store(source_id="src99") | ||
| 145 | + mock_emb.embed_documents.return_value = [[0.1]] | ||
| 146 | + mock_cursor.fetchall.return_value = [(1,)] | ||
| 147 | + | ||
| 148 | + store.add_texts(["text1"]) | ||
| 149 | + | ||
| 150 | + call_args = store._pg_extras.execute_values.call_args | ||
| 151 | + rows = call_args[0][2] | ||
| 152 | + assert rows[0][3] == "src99" | ||
| 153 | + | ||
| 154 | + | ||
| 155 | + | ||
| 156 | +class TestOpenGaussDataVecStoreDeleteIndex: | ||
| 157 | + def test_delete_index_called_with_source_id(self): | ||
| 158 | + store, _, mock_cursor, _ = _make_store(source_id="src123") | ||
| 159 | + | ||
| 160 | + store.delete_index() | ||
| 161 | + | ||
| 162 | + mock_cursor.execute.assert_called_once() | ||
| 163 | + params = mock_cursor.execute.call_args[0][1] | ||
| 164 | + assert params == ("src123",) | ||
| 165 | + | ||
| 166 | + | ||
| 167 | + | ||
| 168 | +class TestOpenGaussDataVecStoreSaveLocal: | ||
| 169 | + def test_save_local_is_noop(self): | ||
| 170 | + store, _, _, _ = _make_store() | ||
| 171 | + assert store.save_local() is None | ||
| 172 | + | ||
| 173 | + | ||
| 174 | + | ||
| 175 | +class TestOpenGaussDataVecStoreGetChunks: | ||
| 176 | + def test_get_chunks(self): | ||
| 177 | + store, _, mock_cursor, _ = _make_store() | ||
| 178 | + mock_cursor.fetchall.return_value = [ | ||
| 179 | + (1, "text1", {"key": "val"}), | ||
| 180 | + (2, "text2", None), | ||
| 181 | + ] | ||
| 182 | + | ||
| 183 | + chunks = store.get_chunks() | ||
| 184 | + assert len(chunks) == 2 | ||
| 185 | + assert chunks[0] == {"doc_id": "1", "text": "text1", "metadata": {"key": "val"}} | ||
| 186 | + assert chunks[1] == {"doc_id": "2", "text": "text2", "metadata": {}} | ||
| 187 | + | ||
| 188 | + def test_get_chunks_returns_empty_on_error(self): | ||
| 189 | + store, _, mock_cursor, _ = _make_store() | ||
| 190 | + mock_cursor.execute.side_effect = Exception("fail") | ||
| 191 | + | ||
| 192 | + assert store.get_chunks() == [] | ||
| 193 | + | ||
| 194 | + def test_get_chunks_filters_by_source_id(self): | ||
| 195 | + store, _, mock_cursor, _ = _make_store(source_id="src7") | ||
| 196 | + mock_cursor.fetchall.return_value = [] | ||
| 197 | + | ||
| 198 | + store.get_chunks() | ||
| 199 | + | ||
| 200 | + params = mock_cursor.execute.call_args[0][1] | ||
| 201 | + assert params == ("src7",) | ||
| 202 | + | ||
| 203 | + | ||
| 204 | + | ||
| 205 | +class TestOpenGaussDataVecStoreAddChunk: | ||
| 206 | + def test_add_chunk_delegates_to_add_texts(self): | ||
| 207 | + store, _, _, _ = _make_store() | ||
| 208 | + store.add_texts = Mock(return_value=["42"]) | ||
| 209 | + | ||
| 210 | + chunk_id = store.add_chunk("hello", metadata={"key": "val"}) | ||
| 211 | + | ||
| 212 | + assert chunk_id == "42" | ||
| 213 | + store.add_texts.assert_called_once_with(["hello"], [{"key": "val"}]) | ||
| 214 | + | ||
| 215 | + def test_add_chunk_default_metadata(self): | ||
| 216 | + store, _, _, _ = _make_store() | ||
| 217 | + store.add_texts = Mock(return_value=["1"]) | ||
| 218 | + | ||
| 219 | + store.add_chunk("text") | ||
| 220 | + | ||
| 221 | + store.add_texts.assert_called_once_with(["text"], [{}]) | ||
| 222 | + | ||
| 223 | + def test_add_chunk_raises_on_empty_result(self): | ||
| 224 | + store, _, _, _ = _make_store() | ||
| 225 | + store.add_texts = Mock(return_value=[]) | ||
| 226 | + | ||
| 227 | + with pytest.raises(IndexError): | ||
| 228 | + store.add_chunk("text") | ||
| 229 | + | ||
| 230 | + | ||
| 231 | + | ||
| 232 | +class TestOpenGaussDataVecStoreDeleteChunk: | ||
| 233 | + def test_delete_chunk_success(self): | ||
| 234 | + store, _, mock_cursor, _ = _make_store() | ||
| 235 | + mock_cursor.rowcount = 1 | ||
| 236 | + | ||
| 237 | + result = store.delete_chunk("42") | ||
| 238 | + assert result is True | ||
| 239 | + | ||
| 240 | + def test_delete_chunk_not_found(self): | ||
| 241 | + store, _, mock_cursor, _ = _make_store() | ||
| 242 | + mock_cursor.rowcount = 0 | ||
| 243 | + | ||
| 244 | + result = store.delete_chunk("999") | ||
| 245 | + assert result is False | ||
| 246 | + | ||
| 247 | + def test_delete_chunk_returns_false_on_error(self): | ||
| 248 | + store, _, mock_cursor, _ = _make_store() | ||
| 249 | + mock_cursor.execute.side_effect = Exception("fail") | ||
| 250 | + | ||
| 251 | + result = store.delete_chunk("42") | ||
| 252 | + assert result is False | ||
| 253 | + | ||
| 254 | + def test_delete_chunk_passes_int_id_and_source_id(self): | ||
| 255 | + store, _, mock_cursor, _ = _make_store(source_id="src5") | ||
| 256 | + mock_cursor.rowcount = 1 | ||
| 257 | + | ||
| 258 | + store.delete_chunk("42") | ||
| 259 | + | ||
| 260 | + params = mock_cursor.execute.call_args[0][1] | ||
| 261 | + assert params == (42, "src5") | ||
| 262 | + | ||
| 263 | + | ||
| 264 | + | ||
| 265 | +class TestOpenGaussDataVecStoreEnsureTable: | ||
| 266 | + def _build_with_table(self, mock_emb): | ||
| 267 | + """Build a store WITHOUT patching _ensure_table_exists.""" | ||
| 268 | + mock_conn = MagicMock() | ||
| 269 | + mock_cursor = MagicMock() | ||
| 270 | + mock_conn.cursor.return_value.__enter__ = Mock(return_value=mock_cursor) | ||
| 271 | + mock_conn.cursor.return_value.__exit__ = Mock(return_value=False) | ||
| 272 | + | ||
| 273 | + with patch.object(OpenGaussDataVecStore, "_load_driver"), \ | ||
| 274 | + patch.object(OpenGaussDataVecStore, "_psycopg2", MagicMock()), \ | ||
| 275 | + patch.object(OpenGaussDataVecStore, "_sql", MagicMock()), \ | ||
| 276 | + patch.object(OpenGaussDataVecStore, "_pg_extras", MagicMock()), \ | ||
| 277 | + patch( | ||
| 278 | + "application.vectorstore.base.BaseVectorStore._get_embeddings", | ||
| 279 | + return_value=mock_emb, | ||
| 280 | + ), \ | ||
| 281 | + patch( | ||
| 282 | + "application.vectorstore.opengauss_datavec.settings" | ||
| 283 | + ) as mock_settings, \ | ||
| 284 | + patch.object( | ||
| 285 | + OpenGaussDataVecStore, "_get_connection", return_value=mock_conn | ||
| 286 | + ): | ||
| 287 | + | ||
| 288 | + mock_settings.EMBEDDINGS_NAME = "test_model" | ||
| 289 | + mock_settings.OPENGAUSS_CONNECTION_STRING = "host=localhost dbname=test" | ||
| 290 | + | ||
| 291 | + store = OpenGaussDataVecStore(source_id="test", embeddings_key="key") | ||
| 292 | + | ||
| 293 | + return store, mock_cursor | ||
| 294 | + | ||
| 295 | + def test_ensure_table_executes_create_statements(self): | ||
| 296 | + mock_emb = Mock(dimension=768, embed_query=Mock(return_value=[0.1, 0.2, 0.3])) | ||
| 297 | + _, mock_cursor = self._build_with_table(mock_emb) | ||
| 298 | + | ||
| 299 | + # CREATE TABLE + 2 indexes (ivfflat + source_id) = 3 execute calls | ||
| 300 | + assert mock_cursor.execute.call_count == 3 | ||
| 301 | + | ||
| 302 | + def test_uses_actual_dimension_when_declared_is_wrong(self): | ||
| 303 | + mock_emb = Mock(dimension=768, embed_query=Mock(return_value=[0.1] * 1536)) | ||
| 304 | + store, _ = self._build_with_table(mock_emb) | ||
| 305 | + | ||
| 306 | + assert store._embedding_dimension == 1536 | ||
| 307 | + assert mock_emb.dimension == 1536 | ||
| 308 | + | ||
| 309 | + def test_probes_when_dimension_missing(self): | ||
| 310 | + mock_emb = Mock(spec=["embed_query"]) | ||
| 311 | + mock_emb.embed_query = Mock(return_value=[0.1, 0.2, 0.3, 0.4]) | ||
| 312 | + store, _ = self._build_with_table(mock_emb) | ||
| 313 | + | ||
| 314 | + assert store._embedding_dimension == 4 | ||
| @@ -0,0 +1,26 @@ | |||
| 1 | +from application.vectorstore.faiss import FaissStore | ||
| 2 | +from application.vectorstore.elasticsearch import ElasticsearchStore | ||
| 3 | +from application.vectorstore.milvus import MilvusStore | ||
| 4 | +from application.vectorstore.mongodb import MongoDBVectorStore | ||
| 5 | +from application.vectorstore.qdrant import QdrantStore | ||
| 6 | +from application.vectorstore.pgvector import PGVectorStore | ||
| 7 | +from application.vectorstore.opengauss_datavec import OpenGaussDataVecStore | ||
| 8 | + | ||
| 9 | + | ||
| 10 | +class VectorCreator: | ||
| 11 | + vectorstores = { | ||
| 12 | + "faiss": FaissStore, | ||
| 13 | + "elasticsearch": ElasticsearchStore, | ||
| 14 | + "mongodb": MongoDBVectorStore, | ||
| 15 | + "qdrant": QdrantStore, | ||
| 16 | + "milvus": MilvusStore, | ||
| 17 | + "pgvector": PGVectorStore, | ||
| 18 | + "opengauss_datavec": OpenGaussDataVecStore, | ||
| 19 | + } | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + def create_vectorstore(cls, type, *args, **kwargs): | ||
| 23 | + vectorstore_class = cls.vectorstores.get(type.lower()) | ||
| 24 | + if not vectorstore_class: | ||
| 25 | + raise ValueError(f"No vectorstore class found for type {type}") | ||
| 26 | + return vectorstore_class(*args, **kwargs) | ||
| @@ -0,0 +1,64 @@ | |||
| 1 | +## 部署 DocsGPT | ||
| 2 | + | ||
| 3 | +本教程将引导你在本地部署 DocsGPT,并将向量存储后端替换为 openGauss DataVec。 | ||
| 4 | + | ||
| 5 | +整个过程分为获取源码、安装依赖、配置环境变量、启动服务四个步骤。 | ||
| 6 | + | ||
| 7 | +### 1. 获取源码 | ||
| 8 | + | ||
| 9 | +首先克隆 DocsGPT 官方仓库: | ||
| 10 | + | ||
| 11 | +```bash | ||
| 12 | +git clone https://github.com/arc53/DocsGPT.git | ||
| 13 | +cd DocsGPT | ||
| 14 | +``` | ||
| 15 | + | ||
| 16 | +克隆完成后,需要将本项目中的相关文件添加或覆盖到 DocsGPT 目录中。涉及的文件如下: | ||
| 17 | + | ||
| 18 | +- **新增文件**:`application/vectorstore/opengauss_datavec.py`、`tests/vectorstore/test_opengauss_datavec.py` | ||
| 19 | +- **修改文件**:`application/core/settings.py`(新增 `OPENGAUSS_CONNECTION_STRING` 配置项)、`application/vectorstore/vector_creator.py`(注册 openGauss DataVec 存储类型) | ||
| 20 | + | ||
| 21 | +如果不想手动复制文件,也可以直接克隆已经集成好的仓库,省去手动替换的步骤: | ||
| 22 | + | ||
| 23 | +```bash | ||
| 24 | +git clone git@github.com:fighting-u/DocsGPT.git | ||
| 25 | +cd DocsGPT | ||
| 26 | +``` | ||
| 27 | + | ||
| 28 | +### 2. 安装依赖 | ||
| 29 | + | ||
| 30 | +源码准备好之后,接下来安装 Python 依赖。先安装 DocsGPT 官方所需的依赖包,再额外安装 openGauss 数据库驱动: | ||
| 31 | + | ||
| 32 | +```bash | ||
| 33 | +pip install -r requirements.txt | ||
| 34 | +pip install psycopg2-binary | ||
| 35 | +``` | ||
| 36 | + | ||
| 37 | +### 3. 配置环境变量 | ||
| 38 | + | ||
| 39 | +依赖安装完成后,需要在项目根目录的 `.env` 文件中添加配置信息,以告知 DocsGPT 使用 openGauss DataVec 作为向量存储后端。以下是一份示例配置: | ||
| 40 | + | ||
| 41 | +```bash | ||
| 42 | +LLM_PROVIDER=docsgpt | ||
| 43 | +VITE_API_STREAMING=true | ||
| 44 | +INTERNAL_KEY=60916fdbecc96881b70f1aba3622f0d585bd41c2d161c3efa8ab4c145f4682d8 | ||
| 45 | +VECTOR_STORE=opengauss_datavec | ||
| 46 | +OPENGAUSS_CONNECTION_STRING=host=10.208.23.46 port=8889 dbname=postgres user=og password=OG@1234 options='-c search_path=public' | ||
| 47 | +EMBEDDINGS_NAME=openai/text-embedding-3-small | ||
| 48 | +EMBEDDINGS_BASE_URL=https://openrouter.ai/api | ||
| 49 | +EMBEDDINGS_KEY=your-key | ||
| 50 | +``` | ||
| 51 | + | ||
| 52 | +### 4. 启动服务 | ||
| 53 | + | ||
| 54 | +一切就绪后,运行启动脚本即可拉起所有服务: | ||
| 55 | + | ||
| 56 | +```bash | ||
| 57 | +./setup.sh | ||
| 58 | +``` | ||
| 59 | + | ||
| 60 | +脚本执行完毕后,在浏览器中打开 http://localhost:5173/ ,即可看到 DocsGPT 的前端界面。 | ||
| 61 | + | ||
| 62 | +此时上传文档后,文本向量将自动存储到 openGauss DataVec 中,你可以直接开始对话式文档问答。 | ||
| 63 | + | ||
| 64 | + | ||