Spaces:
Runtime error
Runtime error
File size: 4,255 Bytes
6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 559a6fa 6b82347 db75ba3 6b82347 db75ba3 6b82347 e206f1f 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | 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 = "google"
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]:
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):
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):
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() |