Files
tk-factory-services/ai-service/db/vector_store.py
Hyungi Ahn b3012b8320 feat: AI 서비스 및 AI 어시스턴트 전용 페이지 추가
- ai-service: Ollama 기반 AI 서비스 (분류, 시맨틱 검색, RAG Q&A, 패턴 분석)
- AI 어시스턴트 페이지: 채팅형 Q&A, 시맨틱 검색, 패턴 분석, 분류 테스트
- 권한 시스템에 ai_assistant 페이지 등록 (기본 비활성)
- 기존 페이지에 AI 기능 통합 (대시보드, 수신함, 관리함)
- docker-compose, gateway, nginx 설정 업데이트

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-06 09:38:30 +09:00

77 lines
2.2 KiB
Python

import chromadb
from config import settings
class VectorStore:
def __init__(self):
self.client = None
self.collection = None
def initialize(self):
self.client = chromadb.PersistentClient(path=settings.CHROMA_PERSIST_DIR)
self.collection = self.client.get_or_create_collection(
name="qc_issues",
metadata={"hnsw:space": "cosine"},
)
def upsert(
self,
doc_id: str,
document: str,
embedding: list[float],
metadata: dict = None,
):
self.collection.upsert(
ids=[doc_id],
documents=[document],
embeddings=[embedding],
metadatas=[metadata] if metadata else None,
)
def query(
self,
embedding: list[float],
n_results: int = 5,
where: dict = None,
) -> list[dict]:
kwargs = {
"query_embeddings": [embedding],
"n_results": n_results,
"include": ["documents", "metadatas", "distances"],
}
if where:
kwargs["where"] = where
try:
results = self.collection.query(**kwargs)
except Exception:
return []
items = []
if results and results["ids"] and results["ids"][0]:
for i, doc_id in enumerate(results["ids"][0]):
item = {
"id": doc_id,
"document": results["documents"][0][i] if results["documents"] else "",
"distance": results["distances"][0][i] if results["distances"] else 0,
"metadata": results["metadatas"][0][i] if results["metadatas"] else {},
}
# cosine distance → similarity
item["similarity"] = round(1 - item["distance"], 4)
items.append(item)
return items
def delete(self, doc_id: str):
self.collection.delete(ids=[doc_id])
def count(self) -> int:
return self.collection.count()
def stats(self) -> dict:
return {
"total_documents": self.count(),
"collection_name": "qc_issues",
}
vector_store = VectorStore()