Initial commit for deployment
This commit is contained in:
Iyeoluwa Akinrinola
@@ -0,0 +1,261 @@
|
||||
"""
|
||||
Service for document processing and chunking.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import uuid
|
||||
import requests
|
||||
import base64
|
||||
from typing import List, Dict, Any, Optional
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
|
||||
from ai_service.config import config
|
||||
|
||||
class DocumentService:
|
||||
"""Service for document processing and chunking."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the document service."""
|
||||
self.chunk_size = config.CHUNK_SIZE
|
||||
self.chunk_overlap = config.CHUNK_OVERLAP
|
||||
self.text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=self.chunk_size,
|
||||
chunk_overlap=self.chunk_overlap,
|
||||
length_function=len
|
||||
)
|
||||
|
||||
# OpenWebUI configuration
|
||||
self.openwebui_url = config.OPENWEBUI_URL
|
||||
self.openwebui_api_key = config.OPENWEBUI_API_KEY
|
||||
|
||||
# Ensure data directory exists
|
||||
os.makedirs(os.path.dirname(config.SQLITE_DB_PATH), exist_ok=True)
|
||||
|
||||
# For now, we'll store document metadata in a simple JSON file
|
||||
self.metadata_file = os.path.join(os.path.dirname(config.SQLITE_DB_PATH), 'document_metadata.json')
|
||||
self._load_metadata()
|
||||
|
||||
def _load_metadata(self):
|
||||
"""Load document metadata from file."""
|
||||
if os.path.exists(self.metadata_file):
|
||||
try:
|
||||
with open(self.metadata_file, 'r') as f:
|
||||
self.documents = json.load(f)
|
||||
except Exception as e:
|
||||
print(f"Error loading document metadata: {str(e)}")
|
||||
self.documents = {}
|
||||
else:
|
||||
self.documents = {}
|
||||
|
||||
def _save_metadata(self):
|
||||
"""Save document metadata to file."""
|
||||
try:
|
||||
with open(self.metadata_file, 'w') as f:
|
||||
json.dump(self.documents, f, indent=2)
|
||||
except Exception as e:
|
||||
print(f"Error saving document metadata: {str(e)}")
|
||||
|
||||
def process_document(self, content: str, title: str,
|
||||
description: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None) -> str:
|
||||
"""
|
||||
Process a document for embedding.
|
||||
|
||||
Args:
|
||||
content: Document content.
|
||||
title: Document title.
|
||||
description: Optional document description.
|
||||
metadata: Optional additional metadata.
|
||||
|
||||
Returns:
|
||||
Document ID.
|
||||
"""
|
||||
# Generate a unique ID for the document
|
||||
doc_id = str(uuid.uuid4())
|
||||
|
||||
# Upload the document to OpenWebUI for RAG processing
|
||||
try:
|
||||
# Prepare headers
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self.openwebui_api_key:
|
||||
headers["Authorization"] = f"Bearer {self.openwebui_api_key}"
|
||||
|
||||
# Prepare the document data
|
||||
document_data = {
|
||||
"filename": f"{title}.txt",
|
||||
"content": base64.b64encode(content.encode('utf-8')).decode('utf-8'),
|
||||
"description": description or title
|
||||
}
|
||||
|
||||
# Upload to OpenWebUI
|
||||
response = requests.post(
|
||||
f"{self.openwebui_url}/api/knowledge/upload",
|
||||
headers=headers,
|
||||
json=document_data,
|
||||
timeout=60
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Get the OpenWebUI document ID
|
||||
openwebui_doc_id = result.get('id', '')
|
||||
|
||||
# Store document metadata
|
||||
self.documents[doc_id] = {
|
||||
'id': doc_id,
|
||||
'title': title,
|
||||
'description': description or '',
|
||||
'openwebui_id': openwebui_doc_id,
|
||||
'metadata': metadata or {}
|
||||
}
|
||||
|
||||
# Save metadata to file
|
||||
self._save_metadata()
|
||||
|
||||
return doc_id
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error uploading document to OpenWebUI: {str(e)}")
|
||||
|
||||
# Fall back to local processing if OpenWebUI upload fails
|
||||
print("Falling back to local document processing")
|
||||
|
||||
# Split the document into chunks for local reference
|
||||
chunks = self.text_splitter.split_text(content)
|
||||
|
||||
# Store document metadata
|
||||
self.documents[doc_id] = {
|
||||
'id': doc_id,
|
||||
'title': title,
|
||||
'description': description or '',
|
||||
'chunk_count': len(chunks),
|
||||
'openwebui_upload_failed': True,
|
||||
'metadata': metadata or {}
|
||||
}
|
||||
|
||||
# Save metadata to file
|
||||
self._save_metadata()
|
||||
|
||||
return doc_id
|
||||
|
||||
def get_document(self, doc_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get document metadata.
|
||||
|
||||
Args:
|
||||
doc_id: Document ID.
|
||||
|
||||
Returns:
|
||||
Document metadata if found, None otherwise.
|
||||
"""
|
||||
return self.documents.get(doc_id)
|
||||
|
||||
def get_all_documents(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get all document metadata.
|
||||
|
||||
Returns:
|
||||
List of document metadata.
|
||||
"""
|
||||
# Get documents from local storage
|
||||
local_documents = list(self.documents.values())
|
||||
|
||||
# Try to get documents from OpenWebUI as well
|
||||
try:
|
||||
# Prepare headers
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self.openwebui_api_key:
|
||||
headers["Authorization"] = f"Bearer {self.openwebui_api_key}"
|
||||
|
||||
# Get documents from OpenWebUI
|
||||
response = requests.get(
|
||||
f"{self.openwebui_url}/api/knowledge",
|
||||
headers=headers,
|
||||
timeout=30
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
openwebui_docs = response.json()
|
||||
|
||||
# Update local documents with OpenWebUI information
|
||||
for doc in local_documents:
|
||||
if 'openwebui_id' in doc:
|
||||
for openwebui_doc in openwebui_docs:
|
||||
if openwebui_doc.get('id') == doc['openwebui_id']:
|
||||
doc['openwebui_status'] = 'active'
|
||||
doc['openwebui_info'] = openwebui_doc
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error getting documents from OpenWebUI: {str(e)}")
|
||||
|
||||
return local_documents
|
||||
|
||||
def delete_document(self, doc_id: str) -> bool:
|
||||
"""
|
||||
Delete a document and its chunks.
|
||||
|
||||
Args:
|
||||
doc_id: Document ID.
|
||||
|
||||
Returns:
|
||||
True if deletion was successful, False otherwise.
|
||||
"""
|
||||
if doc_id not in self.documents:
|
||||
return False
|
||||
|
||||
# Check if document was uploaded to OpenWebUI
|
||||
doc = self.documents[doc_id]
|
||||
openwebui_id = doc.get('openwebui_id')
|
||||
|
||||
if openwebui_id:
|
||||
try:
|
||||
# Prepare headers
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self.openwebui_api_key:
|
||||
headers["Authorization"] = f"Bearer {self.openwebui_api_key}"
|
||||
|
||||
# Delete from OpenWebUI
|
||||
response = requests.delete(
|
||||
f"{self.openwebui_url}/api/knowledge/{openwebui_id}",
|
||||
headers=headers,
|
||||
timeout=30
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
print(f"Warning: Failed to delete document from OpenWebUI: {response.text}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error deleting document from OpenWebUI: {str(e)}")
|
||||
|
||||
# Delete document metadata
|
||||
del self.documents[doc_id]
|
||||
|
||||
# Save metadata to file
|
||||
self._save_metadata()
|
||||
|
||||
return True
|
||||
|
||||
def search_documents(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search for documents similar to a query.
|
||||
|
||||
Args:
|
||||
query: Search query.
|
||||
top_k: Number of results to return.
|
||||
|
||||
Returns:
|
||||
List of similar document chunks with their metadata.
|
||||
"""
|
||||
# Note: We don't need to implement this method anymore since
|
||||
# RAG is handled directly by OpenWebUI when use_rag=True in the model service
|
||||
|
||||
# Return empty results - this is just a placeholder
|
||||
# The actual RAG functionality is in the model_service.generate_response method
|
||||
return []
|
||||
|
||||
|
||||
# Create a singleton instance
|
||||
document_service = DocumentService()
|
||||
@@ -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