Spaces:
Sleeping
Sleeping
helloguddn (PC helloguddn)
feat: 의존성, chat 엔드포인트, gemini function calling, 이미지 추출, Maps 연동, 노동법 RAG(google.text-embedding-004)
f4bf832 | import os | |
| import json | |
| import chromadb | |
| from chromadb.utils.embedding_functions import GoogleGenerativeAiEmbeddingFunction | |
| _collection: chromadb.Collection | None = None | |
| _DATA_PATH = os.path.join(os.path.dirname(__file__), "../data/labor_law_chunks.json") | |
| def _get_collection() -> chromadb.Collection: | |
| global _collection | |
| if _collection is not None: | |
| return _collection | |
| db_path = os.getenv("VECTOR_DB_PATH", "./data/chromadb") | |
| client = chromadb.PersistentClient(path=db_path) | |
| embedding_fn = GoogleGenerativeAiEmbeddingFunction( | |
| api_key=os.getenv("GEMINI_API_KEY", ""), | |
| model_name="models/text-embedding-004", | |
| ) | |
| _collection = client.get_or_create_collection( | |
| name="labor_law", | |
| embedding_function=embedding_fn, | |
| ) | |
| if _collection.count() == 0: | |
| _seed_collection(_collection) | |
| return _collection | |
| def _seed_collection(col: chromadb.Collection) -> None: | |
| with open(_DATA_PATH, encoding="utf-8") as f: | |
| chunks: list[dict] = json.load(f) | |
| col.add( | |
| ids=[c["id"] for c in chunks], | |
| documents=[c["content"] for c in chunks], | |
| metadatas=[ | |
| { | |
| "topic": c["topic"], | |
| "law": c["law"], | |
| "keywords": ", ".join(c.get("keywords", [])), | |
| } | |
| for c in chunks | |
| ], | |
| ) | |
| def search(query: str, n_results: int = 3) -> list[dict]: | |
| col = _get_collection() | |
| results = col.query(query_texts=[query], n_results=n_results) | |
| docs = results["documents"][0] | |
| metas = results["metadatas"][0] | |
| return [ | |
| {"content": doc, "topic": meta["topic"], "law": meta["law"]} | |
| for doc, meta in zip(docs, metas) | |
| ] | |