Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import logging | |
| from typing import List, Dict, Any | |
| from pinecone import Pinecone, PineconeException | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| class PineconeService: | |
| def __init__(self, batch_size: int = 100): | |
| self.batch_size = batch_size | |
| self._pc = None | |
| self._index = None | |
| self._initialized = False | |
| def initialize(self): | |
| """Lazy init (prevents double import issues)""" | |
| if self._initialized: | |
| return | |
| api_key = os.getenv("PINECONE_API_KEY") | |
| index_name = "index" | |
| if not api_key: | |
| raise Exception("Missing PINECONE_API_KEY") | |
| self._pc = Pinecone(api_key=api_key) | |
| self._index = self._pc.Index(index_name) | |
| self._initialized = True | |
| logger.info("Pinecone initialized once") | |
| def upsert( | |
| self, | |
| namespace: str, | |
| document_id: str, | |
| vectors: List[List[float]], | |
| chunks: List[str], | |
| lang_list: List[str], | |
| tokens_list: List[List[str]] | |
| ) -> Dict[str, Any]: | |
| self.initialize() | |
| if not (len(vectors) == len(chunks) == len(lang_list) == len(tokens_list)): | |
| raise ValueError("Input lists must have same length") | |
| vectors_to_upsert = [] | |
| for i, (vector, chunk, lang, tokens) in enumerate( | |
| zip(vectors, chunks, lang_list, tokens_list), 1 | |
| ): | |
| vectors_to_upsert.append({ | |
| "id": f"{document_id}#chunk{i}", | |
| "values": vector, | |
| "metadata": { | |
| "document_id": document_id, | |
| "chunk_number": i, | |
| "chunk_text": chunk, | |
| "language": lang, | |
| "tokens": json.dumps(tokens) | |
| } | |
| }) | |
| responses = [] | |
| for i in range(0, len(vectors_to_upsert), self.batch_size): | |
| batch = vectors_to_upsert[i:i + self.batch_size] | |
| try: | |
| res = self._index.upsert(namespace=namespace, vectors=batch) | |
| responses.append(res) | |
| except Exception as e: | |
| logger.error(f"Upsert failed: {e}") | |
| raise | |
| return {"batches": len(responses)} | |
| def search(self, namespace: str, vector: List[float], doc_ids: List[str], top_k: int): | |
| self.initialize() | |
| all_matches = [] | |
| for doc_id in doc_ids: | |
| try: | |
| result = self._index.query( | |
| namespace=namespace, | |
| vector=vector, | |
| top_k=top_k, | |
| filter={"document_id": {"$eq": doc_id}}, | |
| include_metadata=True | |
| ) | |
| matches = result.get("matches", []) | |
| all_matches.extend(matches) | |
| except PineconeException as e: | |
| logger.error(f"Pinecone error: {e}") | |
| except Exception as e: | |
| logger.error(f"Unexpected error: {e}") | |
| chunk_texts, chunk_ids, tokens_list, lang_list = [], [], [], [] | |
| for match in all_matches: | |
| metadata = match.get("metadata", {}) | |
| chunk_texts.append(metadata.get("chunk_text", "")) | |
| chunk_ids.append(match.get("id", "")) | |
| lang_list.append(metadata.get("language", "unknown")) | |
| tokens_str = metadata.get("tokens", "") | |
| try: | |
| tokens = json.loads(tokens_str) if tokens_str else [] | |
| except: | |
| tokens = [] | |
| tokens_list.append(tokens) | |
| return { | |
| "chunk_texts": chunk_texts, | |
| "chunk_ids": chunk_ids, | |
| "tokens_list": tokens_list, | |
| "lang_list": lang_list | |
| } | |
| def delete(self, namespace: str, document_id: str): | |
| self.initialize() | |
| try: | |
| self._index.delete( | |
| namespace=namespace, | |
| filter={"document_id": {"$eq": document_id}} | |
| ) | |
| except Exception as e: | |
| logger.error(f"Delete failed: {e}") | |
| raise | |
| pinecone_service = PineconeService() |