Spaces:
Sleeping
Sleeping
File size: 3,304 Bytes
9002555 |
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 |
from llama_index.core import VectorStoreIndex
from llama_index.core import StorageContext
from pinecone import Pinecone, ServerlessSpec
from llama_index.llms.openai import OpenAI
from llama_index.vector_stores.pinecone import PineconeVectorStore
from fastapi import HTTPException, status
from config import PINECONE_CONFIG
import os
import json
class IndexManager:
def __init__(self):
self.vector_index = None
self.index_name = "summarizer-semantic-index"
def _get_pinecone_client(self):
"""Initialize and return the Pinecone client."""
# api_key = os.getenv("PINECONE_API_KEY")
api_key = PINECONE_CONFIG.PINECONE_API_KEY
if not api_key:
raise ValueError(
"Pinecone API key is missing. Please set it in environment variables."
)
return Pinecone(api_key=api_key)
def _create_pinecone_index(self, client):
"""Create Pinecone index if it doesn't already exist."""
if self.index_name not in client.list_indexes().names():
client.create_index(
name=self.index_name,
dimension=1536,
metric="cosine",
spec=ServerlessSpec(cloud="aws", region="us-east-1"),
)
return client.Index(self.index_name)
def _initialize_vector_store(self, pinecone_index):
"""Initialize and return the vector store with the Pinecone index."""
vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
return StorageContext.from_defaults(vector_store=vector_store)
def build_indexes(self, nodes):
"""Build vector and tree indexes from nodes."""
try:
client = self._get_pinecone_client()
pinecone_index = self._create_pinecone_index(client)
storage_context = self._initialize_vector_store(pinecone_index)
self.vector_index = VectorStoreIndex(nodes, storage_context=storage_context)
self.vector_index.set_index_id("vector")
print(f"Vector Index ID: {self.vector_index.index_id}")
print("Vector Index created successfully.")
response = {
"status": "success",
"message": "Existing Vector Index loaded successfully.",
}
return json.dumps(response)
except HTTPException as http_exc:
raise http_exc # Re-raise HTTPExceptions to ensure FastAPI handles them
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error loading existing indexes: {str(e)}"
)
def load_existing_indexes(self):
"""Load existing indexes from Pinecone."""
try:
client = self._get_pinecone_client()
pinecone_index = client.Index(self.index_name)
print(pinecone_index.describe_index_stats())
vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
retriever = VectorStoreIndex.from_vector_store(vector_store)
print("Existing Vector Index loaded successfully.")
return retriever
except Exception as e:
print(f"Error loading existing indexes: {e}")
raise |