Initial commit for deployment
This commit is contained in:
Iyeoluwa Akinrinola
@@ -0,0 +1,214 @@
|
||||
"""
|
||||
Service for generating and managing embeddings.
|
||||
"""
|
||||
|
||||
import os
|
||||
import random
|
||||
import pinecone
|
||||
import numpy as np
|
||||
from typing import List, Dict, Any, Optional, Union
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
from ai_service.config import config
|
||||
|
||||
class EmbeddingService:
|
||||
"""Service for generating and managing embeddings."""
|
||||
|
||||
def __init__(self, use_mock=True): # Default to mock implementation
|
||||
"""Initialize the embedding service."""
|
||||
self.use_mock = use_mock
|
||||
|
||||
if not self.use_mock:
|
||||
# Use a smaller model for testing
|
||||
self.model_name = "paraphrase-MiniLM-L3-v2" # Smaller model than the default
|
||||
try:
|
||||
self.model = SentenceTransformer(self.model_name)
|
||||
print(f"Loaded embedding model: {self.model_name}")
|
||||
except Exception as e:
|
||||
print(f"Error loading embedding model: {str(e)}")
|
||||
self.use_mock = True
|
||||
print("Falling back to mock implementation")
|
||||
else:
|
||||
print("Using mock embedding implementation")
|
||||
self.model_name = "mock-model"
|
||||
self.model = None
|
||||
|
||||
self._initialize_pinecone()
|
||||
|
||||
def _initialize_pinecone(self):
|
||||
"""Initialize Pinecone client."""
|
||||
if not config.PINECONE_API_KEY or not config.PINECONE_ENVIRONMENT:
|
||||
print("Warning: Pinecone API key or environment not set. Vector storage will not be available.")
|
||||
self.index = None
|
||||
return
|
||||
|
||||
try:
|
||||
pinecone.init(
|
||||
api_key=config.PINECONE_API_KEY,
|
||||
environment=config.PINECONE_ENVIRONMENT
|
||||
)
|
||||
|
||||
# Check if index exists, create if it doesn't
|
||||
if config.PINECONE_INDEX_NAME not in pinecone.list_indexes():
|
||||
pinecone.create_index(
|
||||
name=config.PINECONE_INDEX_NAME,
|
||||
dimension=self.model.get_sentence_embedding_dimension(),
|
||||
metric="cosine"
|
||||
)
|
||||
|
||||
self.index = pinecone.Index(config.PINECONE_INDEX_NAME)
|
||||
print(f"Connected to Pinecone index: {config.PINECONE_INDEX_NAME}")
|
||||
except Exception as e:
|
||||
print(f"Error connecting to Pinecone: {str(e)}")
|
||||
self.index = None
|
||||
|
||||
def generate_embedding(self, text: str) -> List[float]:
|
||||
"""
|
||||
Generate an embedding for a text.
|
||||
|
||||
Args:
|
||||
text: Text to embed.
|
||||
|
||||
Returns:
|
||||
Embedding vector.
|
||||
"""
|
||||
if self.use_mock:
|
||||
# Generate a mock embedding vector (384 dimensions for consistency)
|
||||
return [random.random() for _ in range(384)]
|
||||
|
||||
embedding = self.model.encode(text)
|
||||
return embedding.tolist()
|
||||
|
||||
def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Generate embeddings for multiple texts.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed.
|
||||
|
||||
Returns:
|
||||
List of embedding vectors.
|
||||
"""
|
||||
if self.use_mock:
|
||||
# Generate mock embedding vectors
|
||||
return [[random.random() for _ in range(384)] for _ in texts]
|
||||
|
||||
embeddings = self.model.encode(texts)
|
||||
return embeddings.tolist()
|
||||
|
||||
def store_embeddings(self, ids: List[str], embeddings: List[List[float]],
|
||||
metadata: Optional[List[Dict[str, Any]]] = None) -> bool:
|
||||
"""
|
||||
Store embeddings in Pinecone.
|
||||
|
||||
Args:
|
||||
ids: List of IDs for the embeddings.
|
||||
embeddings: List of embedding vectors.
|
||||
metadata: Optional list of metadata dictionaries.
|
||||
|
||||
Returns:
|
||||
True if storage was successful, False otherwise.
|
||||
"""
|
||||
if self.use_mock:
|
||||
print(f"Mock: Stored {len(ids)} embeddings")
|
||||
return True
|
||||
|
||||
if self.index is None:
|
||||
print("Warning: Pinecone index not available. Embeddings not stored.")
|
||||
return False
|
||||
|
||||
if metadata is None:
|
||||
metadata = [{} for _ in ids]
|
||||
|
||||
vectors = [
|
||||
(id, embedding, meta)
|
||||
for id, embedding, meta in zip(ids, embeddings, metadata)
|
||||
]
|
||||
|
||||
try:
|
||||
self.index.upsert(vectors=vectors)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error storing embeddings in Pinecone: {str(e)}")
|
||||
return False
|
||||
|
||||
def search_similar(self, query_embedding: List[float], top_k: int = 5) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search for similar embeddings in Pinecone.
|
||||
|
||||
Args:
|
||||
query_embedding: Query embedding vector.
|
||||
top_k: Number of results to return.
|
||||
|
||||
Returns:
|
||||
List of similar items with their metadata.
|
||||
"""
|
||||
if self.use_mock:
|
||||
# Generate mock search results
|
||||
print(f"Mock: Searching for similar embeddings (top_k={top_k})")
|
||||
mock_results = []
|
||||
for i in range(min(top_k, 3)): # Return at most 3 mock results
|
||||
mock_results.append({
|
||||
'id': f"mock_doc_{i}",
|
||||
'score': 0.9 - (i * 0.1), # Decreasing similarity scores
|
||||
'metadata': {
|
||||
'document_id': f"mock_doc_{i}",
|
||||
'chunk_index': i,
|
||||
'title': f"Mock Document {i}",
|
||||
'description': f"This is a mock document {i}",
|
||||
'chunk_text': f"This is the content of mock document {i}..."
|
||||
}
|
||||
})
|
||||
return mock_results
|
||||
|
||||
if self.index is None:
|
||||
print("Warning: Pinecone index not available. Search not performed.")
|
||||
return []
|
||||
|
||||
try:
|
||||
results = self.index.query(
|
||||
vector=query_embedding,
|
||||
top_k=top_k,
|
||||
include_metadata=True
|
||||
)
|
||||
|
||||
return [
|
||||
{
|
||||
'id': match['id'],
|
||||
'score': match['score'],
|
||||
'metadata': match.get('metadata', {})
|
||||
}
|
||||
for match in results.get('matches', [])
|
||||
]
|
||||
except Exception as e:
|
||||
print(f"Error searching in Pinecone: {str(e)}")
|
||||
return []
|
||||
|
||||
def delete_embeddings(self, ids: List[str]) -> bool:
|
||||
"""
|
||||
Delete embeddings from Pinecone.
|
||||
|
||||
Args:
|
||||
ids: List of IDs to delete.
|
||||
|
||||
Returns:
|
||||
True if deletion was successful, False otherwise.
|
||||
"""
|
||||
if self.use_mock:
|
||||
print(f"Mock: Deleted {len(ids)} embeddings")
|
||||
return True
|
||||
|
||||
if self.index is None:
|
||||
print("Warning: Pinecone index not available. Deletion not performed.")
|
||||
return False
|
||||
|
||||
try:
|
||||
self.index.delete(ids=ids)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error deleting embeddings from Pinecone: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
# Create a singleton instance
|
||||
embedding_service = EmbeddingService()
|
||||
Reference in New Issue
Block a user