RAG implementation, content moderation, prompt classification, new LLM chain, document storage
This commit is contained in:
@@ -0,0 +1,145 @@
|
||||
import os
|
||||
import logging
|
||||
from typing import Optional, Tuple
|
||||
from PIL import Image
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline, DPMSolverSinglestepScheduler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ImageGenerationService:
|
||||
"""
|
||||
Service for text-to-image generation using Stable Diffusion.
|
||||
Uses singleton pattern to maintain loaded model in memory.
|
||||
"""
|
||||
|
||||
_instance = None
|
||||
_model_loaded = False
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._initialize()
|
||||
return cls._instance
|
||||
|
||||
def _initialize(self):
|
||||
"""Initialize the service with default settings"""
|
||||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
self.model_id = "stabilityai/stable-diffusion-2-1"
|
||||
self.pipeline = None
|
||||
self.default_params = {
|
||||
"num_inference_steps": 25,
|
||||
"guidance_scale": 7.5,
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
}
|
||||
|
||||
def load_model(self):
|
||||
"""Load the Stable Diffusion model"""
|
||||
if self._model_loaded:
|
||||
return
|
||||
|
||||
try:
|
||||
logger.info(f"Loading Stable Diffusion model on {self.device}...")
|
||||
|
||||
# Use DPMSolver for faster inference
|
||||
self.pipeline = StableDiffusionPipeline.from_pretrained(
|
||||
self.model_id,
|
||||
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
|
||||
)
|
||||
self.pipeline.scheduler = DPMSolverSinglestepScheduler.from_config(
|
||||
self.pipeline.scheduler.config
|
||||
)
|
||||
self.pipeline = self.pipeline.to(self.device)
|
||||
|
||||
# Optimizations
|
||||
if self.device == "cuda":
|
||||
self.pipeline.enable_attention_slicing()
|
||||
self.pipeline.enable_xformers_memory_efficient_attention()
|
||||
|
||||
self._model_loaded = True
|
||||
logger.info("Model loaded successfully")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load model: {str(e)}")
|
||||
raise RuntimeError(f"Model loading failed: {str(e)}")
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
prompt: str,
|
||||
negative_prompt: Optional[str] = None,
|
||||
output_path: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Tuple[Image.Image, dict]:
|
||||
"""
|
||||
Generate image from text prompt.
|
||||
|
||||
Args:
|
||||
prompt: Text prompt for image generation
|
||||
negative_prompt: Text for things to avoid in generation
|
||||
output_path: Optional path to save the image
|
||||
**kwargs: Generation parameters (overrides defaults)
|
||||
|
||||
Returns:
|
||||
Tuple of (PIL.Image, generation_parameters)
|
||||
"""
|
||||
if not self._model_loaded:
|
||||
self.load_model()
|
||||
|
||||
# Merge default params with overrides
|
||||
params = {**self.default_params, **kwargs}
|
||||
|
||||
try:
|
||||
logger.info(f"Generating image with prompt: {prompt[:50]}...")
|
||||
|
||||
with torch.inference_mode():
|
||||
result = self.pipeline(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
**params
|
||||
)
|
||||
|
||||
image = result.images[0]
|
||||
|
||||
if output_path:
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
image.save(output_path)
|
||||
logger.info(f"Image saved to {output_path}")
|
||||
|
||||
return image, params
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Image generation failed: {str(e)}")
|
||||
raise RuntimeError(f"Image generation failed: {str(e)}")
|
||||
|
||||
|
||||
class AsyncImageGenerationService:
|
||||
"""
|
||||
Asynchronous wrapper for image generation service.
|
||||
Runs the synchronous service in a thread pool.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.sync_service = ImageGenerationService()
|
||||
|
||||
async def generate_image(
|
||||
self,
|
||||
prompt: str,
|
||||
negative_prompt: Optional[str] = None,
|
||||
output_path: Optional[str] = None,
|
||||
**kwargs
|
||||
) -> Tuple[Image.Image, dict]:
|
||||
"""Async version of generate_image"""
|
||||
import asyncio
|
||||
from functools import partial
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
func = partial(
|
||||
self.sync_service.generate_image,
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
output_path=output_path,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
return await loop.run_in_executor(None, func)
|
||||
@@ -0,0 +1,138 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import AsyncGenerator, Generator, Optional
|
||||
|
||||
from langchain_community.llms import Ollama
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
|
||||
from chat_backend.models import Conversation, Prompt
|
||||
|
||||
class LLMService(ABC):
|
||||
"""Abstract base class for LLM conversation services."""
|
||||
|
||||
def __init__(self):
|
||||
self.llm = Ollama(
|
||||
model="llama3.2",
|
||||
temperature=0.7,
|
||||
top_k=50,
|
||||
top_p=0.9,
|
||||
repeat_penalty=1.1,
|
||||
num_ctx=4096
|
||||
)
|
||||
self.output_parser = StrOutputParser()
|
||||
|
||||
@abstractmethod
|
||||
def generate_response(self, conversation: Conversation, query: str, **kwargs):
|
||||
"""Generate a response to a query within a conversation context."""
|
||||
pass
|
||||
|
||||
def _format_history(self, conversation: Conversation) -> str:
|
||||
"""Format conversation history for the prompt."""
|
||||
prompts = Prompt.objects.filter(conversation=conversation).order_by('created_at')
|
||||
return "\n".join(
|
||||
f"{'User' if prompt.is_user else 'AI'}: {prompt.text}"
|
||||
for prompt in prompts
|
||||
)
|
||||
|
||||
|
||||
class SyncLLMService(LLMService):
|
||||
"""Synchronous LLM conversation service."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._setup_chain()
|
||||
|
||||
def _setup_chain(self):
|
||||
"""Setup the conversation chain."""
|
||||
template = """Continue the conversation based on the following history:
|
||||
|
||||
{history}
|
||||
|
||||
Latest message: {query}
|
||||
|
||||
Response:"""
|
||||
self.prompt = ChatPromptTemplate.from_template(template)
|
||||
|
||||
self.conversation_chain = (
|
||||
{
|
||||
"history": lambda x: self._format_history(x["conversation"]),
|
||||
"query": lambda x: x["query"]
|
||||
}
|
||||
| self.prompt
|
||||
| self.llm
|
||||
| self.output_parser
|
||||
)
|
||||
|
||||
def generate_response(self, conversation: Conversation, query: str, **kwargs) -> Generator[str, None, None]:
|
||||
"""Generate response with streaming support."""
|
||||
chain_input = {
|
||||
"query": query,
|
||||
"conversation": conversation
|
||||
}
|
||||
|
||||
for chunk in self.conversation_chain.stream(chain_input):
|
||||
yield chunk
|
||||
|
||||
|
||||
class AsyncLLMService(LLMService):
|
||||
"""Asynchronous LLM conversation service."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._setup_chain()
|
||||
|
||||
def _setup_chain(self):
|
||||
"""Setup the conversation chain."""
|
||||
template = """Continue this conversation while maintaining context by providing a single helpful response.
|
||||
Current context: {context}
|
||||
|
||||
Last 3 messages:
|
||||
{recent_history}
|
||||
|
||||
Latest message: {query}
|
||||
|
||||
Instructions:
|
||||
- Carefully maintain all established context
|
||||
- If referencing previous elements (like stories), preserve all details
|
||||
- When asked to modify something, identify what's being modified
|
||||
|
||||
Response:"""
|
||||
|
||||
self.prompt = ChatPromptTemplate.from_template(template)
|
||||
|
||||
self.conversation_chain = (
|
||||
{
|
||||
"context": lambda x: self._format_history(x["conversation"]),
|
||||
"recent_history": lambda x: self._get_recent_messages(x["conversation"]),
|
||||
"query": lambda x: x["query"]
|
||||
}
|
||||
| self.prompt
|
||||
| self.llm
|
||||
| self.output_parser
|
||||
)
|
||||
|
||||
async def _format_history(self, conversation: Conversation) -> str:
|
||||
"""Async version of format conversation history."""
|
||||
prompts = await Prompt.objects.filter(conversation=conversation).order_by('created_at').alist()
|
||||
return "\n".join(
|
||||
f"{'User' if prompt.is_user else 'AI'}: {prompt.text}"
|
||||
for prompt in prompts
|
||||
)
|
||||
|
||||
async def _get_recent_messages(self, conversation: Conversation) -> str:
|
||||
"""Async version of format conversation history."""
|
||||
prompts = await Prompt.objects.filter(conversation=conversation).order_by('created_at').alist()[-3:]
|
||||
return "\n".join(
|
||||
f"{'User' if prompt.is_user else 'AI'}: {prompt.text}"
|
||||
for prompt in prompts
|
||||
)
|
||||
|
||||
async def generate_response(self, conversation: Conversation, query: str, **kwargs) -> AsyncGenerator[str, None]:
|
||||
"""Generate response with async streaming support."""
|
||||
chain_input = {
|
||||
"query": query,
|
||||
"conversation": conversation
|
||||
}
|
||||
|
||||
async for chunk in self.conversation_chain.astream(chain_input):
|
||||
yield chunk
|
||||
@@ -0,0 +1,79 @@
|
||||
from enum import Enum, auto
|
||||
from typing import Dict, Any
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_community.llms import Ollama
|
||||
|
||||
class ModerationLabel(Enum):
|
||||
NSFW = auto()
|
||||
FINE = auto()
|
||||
|
||||
class ModerationClassifier:
|
||||
"""
|
||||
Classifies prompts as NSFW or FINE (safe) content.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.llm = Ollama(
|
||||
model="llama3.2",
|
||||
temperature=0.1, # Very low for strict moderation
|
||||
top_k=10,
|
||||
num_ctx=2048
|
||||
)
|
||||
|
||||
self.moderation_prompt = ChatPromptTemplate.from_messages([
|
||||
("system", """You are a strict content moderator. Classify the following prompt as either NSFW or FINE.
|
||||
|
||||
NSFW includes:
|
||||
- Sexual content
|
||||
- Violence/gore
|
||||
- Hate speech
|
||||
- Illegal activities
|
||||
- Harassment
|
||||
- Graphic/disturbing content
|
||||
|
||||
FINE includes:
|
||||
- Safe for work topics
|
||||
- General conversation
|
||||
- Professional inquiries
|
||||
- Creative requests (non-explicit)
|
||||
- Technical questions
|
||||
|
||||
Examples:
|
||||
- "How to make a bomb" → NSFW
|
||||
- "Write a love poem" → FINE
|
||||
- "Explicit sex scene" → NSFW
|
||||
- "Python tutorial" → FINE
|
||||
|
||||
Return ONLY "NSFW" or "FINE", nothing else."""),
|
||||
("human", "{prompt}")
|
||||
])
|
||||
|
||||
self.chain = self.moderation_prompt | self.llm
|
||||
|
||||
async def classify_async(self, prompt: str) -> ModerationLabel:
|
||||
"""Asynchronous classification"""
|
||||
try:
|
||||
response = (await self.chain.ainvoke({"prompt": prompt})).strip().upper()
|
||||
return self._parse_response(response)
|
||||
except Exception as e:
|
||||
print(f"Moderation error: {e}")
|
||||
return ModerationLabel.NSFW # Fail-safe to NSFW
|
||||
|
||||
def classify(self, prompt: str) -> ModerationLabel:
|
||||
"""Synchronous classification"""
|
||||
try:
|
||||
response = self.chain.invoke({"prompt": prompt}).strip().upper()
|
||||
return self._parse_response(response)
|
||||
except Exception as e:
|
||||
print(f"Moderation error: {e}")
|
||||
return ModerationLabel.NSFW # Fail-safe to NSFW
|
||||
|
||||
def _parse_response(self, response: str) -> ModerationLabel:
|
||||
"""Convert string response to ModerationLabel enum"""
|
||||
if "NSFW" in response:
|
||||
return ModerationLabel.NSFW
|
||||
return ModerationLabel.FINE # Default to FINE if unclear
|
||||
|
||||
|
||||
# Singleton instance
|
||||
moderation_classifier = ModerationClassifier()
|
||||
@@ -0,0 +1,100 @@
|
||||
from enum import Enum, auto
|
||||
from typing import Dict, Any
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_community.llms import Ollama
|
||||
|
||||
class PromptType(Enum):
|
||||
GENERAL_CHAT = auto()
|
||||
RAG = auto()
|
||||
IMAGE_GENERATION = auto()
|
||||
UNKNOWN = auto()
|
||||
|
||||
class PromptClassifier:
|
||||
"""
|
||||
Classifies user prompts to determine which service should handle them.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.llm = Ollama(
|
||||
model="llama3",
|
||||
temperature=0.3, # Lower temp for more deterministic classification
|
||||
top_k=20,
|
||||
top_p=0.9,
|
||||
num_ctx=4096
|
||||
)
|
||||
|
||||
self.classification_prompt = ChatPromptTemplate.from_messages([
|
||||
("system",
|
||||
"""You are a precision prompt classifier. Strictly categorize prompts into:
|
||||
1. GENERAL_CHAT - Casual conversation, personal questions, or non-specific inquiries
|
||||
2. RAG - ONLY when explicitly requesting document/search-based knowledge
|
||||
3. IMAGE_GENERATION - Specific requests to create/modify images
|
||||
4. UNKNOWN - If none of the above fit
|
||||
|
||||
1. IMAGE_GENERATION - ONLY if:
|
||||
- Explicitly contains: "generate/create/draw/make an image/picture/photo/art/illustration"
|
||||
- Requests visual content creation
|
||||
- Example: "Make a picture of a castle" → IMAGE_GENERATION
|
||||
|
||||
2. RAG - ONLY if:
|
||||
- Explicitly mentions documents/files/data
|
||||
- Uses search terms: "find/search/lookup in [source]"
|
||||
- Example: "What does contracts.pdf say?" → RAG
|
||||
|
||||
3. GENERAL_CHAT - DEFAULT category when:
|
||||
- Doesn't meet above criteria
|
||||
- Conversational/general knowledge
|
||||
- Uncertain cases
|
||||
- Example: "Tell me a joke" → GENERAL_CHAT
|
||||
|
||||
Examples:
|
||||
[Definitely RAG]
|
||||
- "What does the uploaded PDF say about quarterly results?"
|
||||
- "Search our documents for the 2023 marketing strategy"
|
||||
- "Find the contract clause about termination"
|
||||
|
||||
[Definitely GENERAL_CHAT]
|
||||
- "How does photosynthesis work?" (General knowledge)
|
||||
- "Tell me a joke"
|
||||
- "What's your opinion on AI?"
|
||||
|
||||
[Borderline → GENERAL_CHAT]
|
||||
- "What's our company policy on X?" (No doc reference → general)
|
||||
- "Explain quantum computing" (General knowledge)
|
||||
- "Summarize the meeting" (No doc reference)
|
||||
|
||||
Return ONLY the label, no explanations."""),
|
||||
("human", "{prompt}")
|
||||
])
|
||||
|
||||
self.chain = self.classification_prompt | self.llm
|
||||
|
||||
async def classify_async(self, prompt: str) -> PromptType:
|
||||
"""Asynchronously classify the prompt"""
|
||||
try:
|
||||
response = await self.chain.ainvoke({"prompt": prompt})
|
||||
return self._parse_response(response.strip())
|
||||
except Exception as e:
|
||||
print(f"Classification error: {e}")
|
||||
return PromptType.UNKNOWN
|
||||
|
||||
def classify(self, prompt: str) -> PromptType:
|
||||
"""Synchronously classify the prompt"""
|
||||
try:
|
||||
response = self.chain.invoke({"prompt": prompt})
|
||||
return self._parse_response(response.strip())
|
||||
except Exception as e:
|
||||
print(f"Classification error: {e}")
|
||||
return PromptType.UNKNOWN
|
||||
|
||||
def _parse_response(self, response: str) -> PromptType:
|
||||
"""Convert string response to PromptType enum"""
|
||||
response = response.upper()
|
||||
for prompt_type in PromptType:
|
||||
if prompt_type.name in response:
|
||||
return prompt_type
|
||||
return PromptType.UNKNOWN
|
||||
|
||||
|
||||
# Singleton instance for easy access
|
||||
prompt_classifier = PromptClassifier()
|
||||
@@ -0,0 +1,378 @@
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Dict, Any, AsyncGenerator, Generator, Optional
|
||||
from channels.db import database_sync_to_async
|
||||
from langchain_community.embeddings import OllamaEmbeddings
|
||||
from langchain_community.llms import Ollama
|
||||
from langchain_community.vectorstores import Chroma
|
||||
from langchain_core.documents import Document as LangDocument
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.runnables import RunnablePassthrough
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from langchain_community.document_loaders import (
|
||||
PyPDFLoader,
|
||||
Docx2txtLoader,
|
||||
TextLoader,
|
||||
UnstructuredFileLoader
|
||||
)
|
||||
from django.core.files.uploadedfile import UploadedFile
|
||||
from chat_backend.models import Conversation, Prompt, DocumentWorkspace, Document
|
||||
from pathlib import Path
|
||||
|
||||
@database_sync_to_async
|
||||
def get_documents(workspace: DocumentWorkspace | None = None):
|
||||
if workspace:
|
||||
return [doc for doc in Document.objects.filter(workspace=workspace)]
|
||||
else:
|
||||
return [doc for doc in Document.objects.all()]
|
||||
|
||||
|
||||
|
||||
class RAGService(ABC):
|
||||
"""Abstract base class for RAG services."""
|
||||
_instance = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance.__init__()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
self.embedding_model = OllamaEmbeddings(model="llama3.2")
|
||||
self.llm = Ollama(
|
||||
model="llama3.2",
|
||||
temperature=0.7,
|
||||
top_k=50,
|
||||
top_p=0.9,
|
||||
repeat_penalty=1.1,
|
||||
num_ctx=4096
|
||||
)
|
||||
self.text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1000,
|
||||
chunk_overlap=200
|
||||
)
|
||||
self.vector_store = self._initialize_vector_store()
|
||||
|
||||
# Supported file types and their loaders
|
||||
self.loader_mapping = {
|
||||
'.pdf': PyPDFLoader,
|
||||
'.docx': Docx2txtLoader,
|
||||
'.txt': TextLoader,
|
||||
# Fallback for other file types
|
||||
'*': UnstructuredFileLoader,
|
||||
}
|
||||
|
||||
def _initialize_vector_store(self) -> Chroma:
|
||||
"""Initialize and return the Chroma vector store."""
|
||||
persist_directory=f"./chroma_db/"
|
||||
vector_store = Chroma(
|
||||
embedding_function=self.embedding_model,
|
||||
persist_directory=persist_directory
|
||||
)
|
||||
return vector_store
|
||||
|
||||
def clear_vector_store(self):
|
||||
"""Clear all vectors from the store"""
|
||||
self.vector_store.delete_collection()
|
||||
self.vector_store = self._initialize_vector_store()
|
||||
|
||||
def _prepare_documents(self, documents: List[Document]) -> List[Document]:
|
||||
"""Process documents for ingestion into vector store."""
|
||||
docs = []
|
||||
|
||||
for doc in documents:
|
||||
print(f"Processing: {doc.file.name}")
|
||||
loader_class = self._get_file_loader( doc.file.name)
|
||||
loader = loader_class(doc.file)
|
||||
|
||||
|
||||
chunks = self._load_and_split_documents(doc.file.path)
|
||||
if chunks:
|
||||
self.vector_store.add_documents(chunks)
|
||||
self.vector_store.persist()
|
||||
|
||||
|
||||
def ingest_documents(self, workspace: DocumentWorkspace | None = None) -> None:
|
||||
"""Ingest documents from a workspace into the vector store."""
|
||||
print(f"Getting the Document via the workspace: {workspace}")
|
||||
if workspace:
|
||||
documents = [doc for doc in Document.objects.filter(workspace=workspace)]
|
||||
else:
|
||||
documents = [doc for doc in Document.objects.all()]
|
||||
|
||||
print(f"Processing the documents : {documents}")
|
||||
self._prepare_documents(documents)
|
||||
|
||||
|
||||
@abstractmethod
|
||||
def generate_response(self, conversation: Conversation, query: str, **kwargs):
|
||||
"""Generate a response using RAG."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def search_documents(self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4) -> List[Document]:
|
||||
"""Search relevant documents from the vector store."""
|
||||
pass
|
||||
|
||||
def _get_file_loader(self, file_path: str):
|
||||
"""Get appropriate loader for file type"""
|
||||
ext = Path(file_path).suffix.lower()
|
||||
return self.loader_mapping.get(ext, self.loader_mapping['*'])
|
||||
|
||||
def _sanitize_filename(self, filename: str) -> str:
|
||||
"""Sanitize filename for safe storage"""
|
||||
return re.sub(r'[^\w\-_. ]', '_', filename)
|
||||
|
||||
def _save_uploaded_file(self, uploaded_file: UploadedFile, save_dir: str) -> str:
|
||||
"""Save uploaded file to disk"""
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
sanitized_name = self._sanitize_filename(uploaded_file.name)
|
||||
file_path = os.path.join(save_dir, sanitized_name)
|
||||
|
||||
with open(file_path, 'wb+') as destination:
|
||||
for chunk in uploaded_file.chunks():
|
||||
destination.write(chunk)
|
||||
|
||||
return file_path
|
||||
|
||||
def _load_and_split_documents(self, file_path: str, metadata: dict = None) -> List[Document]:
|
||||
"""Load and split documents from file"""
|
||||
loader_class = self._get_file_loader(file_path)
|
||||
loader = loader_class(file_path)
|
||||
|
||||
docs = loader.load()
|
||||
if metadata:
|
||||
for doc in docs:
|
||||
doc.metadata.update(metadata)
|
||||
|
||||
return self.text_splitter.split_documents(docs)
|
||||
|
||||
def add_files_to_store(
|
||||
self,
|
||||
file_tupls: List[UploadedFile], # (file_path, name,workspace_id)
|
||||
workspace_id: str,
|
||||
source: str = "upload",
|
||||
save_dir: str = "data/uploads"
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Process and add uploaded files to vector store
|
||||
|
||||
Args:
|
||||
files: List of Django UploadedFile objects
|
||||
workspace_id: ID of the workspace these belong to
|
||||
source: Source identifier for documents
|
||||
save_dir: Directory to save uploaded files
|
||||
|
||||
Returns:
|
||||
Dictionary with processing results
|
||||
"""
|
||||
results = {
|
||||
'total_added': 0,
|
||||
'failed_files': [],
|
||||
'processed_files': []
|
||||
}
|
||||
|
||||
for file_tuple in file_tupls:
|
||||
try:
|
||||
# Save file to disk
|
||||
|
||||
|
||||
# Prepare metadata
|
||||
metadata = {
|
||||
'source': file_tuple[1],
|
||||
'workspace_id': file_tuple[2],
|
||||
'original_filename': file_tuple[1],
|
||||
'file_path': file_tuple[0],
|
||||
}
|
||||
|
||||
# Load and split documents
|
||||
docs = self._load_and_split_documents(file_path, metadata)
|
||||
|
||||
# Add to vector store
|
||||
if docs:
|
||||
self.vector_store.add_documents(docs)
|
||||
results['total_added'] += len(docs)
|
||||
results['processed_files'].append({
|
||||
'filename': file_tuple[1],
|
||||
'document_count': len(docs)
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
results['failed_files'].append({
|
||||
'filename': file_tuple[1],
|
||||
'error': str(e)
|
||||
})
|
||||
continue
|
||||
|
||||
# Persist changes
|
||||
self.vector_store.persist()
|
||||
return results
|
||||
|
||||
|
||||
class SyncRAGService(RAGService):
|
||||
"""Synchronous RAG service implementation."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._setup_chain()
|
||||
|
||||
def _setup_chain(self):
|
||||
"""Setup the RAG chain."""
|
||||
template = """Answer the question based only on the following context:
|
||||
{context}
|
||||
|
||||
Conversation history:
|
||||
{history}
|
||||
|
||||
Question: {question}
|
||||
"""
|
||||
self.prompt = ChatPromptTemplate.from_template(template)
|
||||
|
||||
self.rag_chain = (
|
||||
{
|
||||
"context": self._retriever_with_history,
|
||||
"history": lambda x: self._format_history(x["conversation"]),
|
||||
"question": lambda x: x["query"]
|
||||
}
|
||||
| self.prompt
|
||||
| self.llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
def _format_history(self, conversation: Conversation) -> str:
|
||||
"""Format conversation history for the prompt."""
|
||||
prompts = Prompt.objects.filter(conversation=conversation).order_by('created_at')
|
||||
return "\n".join(
|
||||
f"{'User' if prompt.is_user else 'AI'}: {prompt.text}"
|
||||
for prompt in prompts
|
||||
)
|
||||
|
||||
def _retriever_with_history(self, input_dict: Dict[str, Any]) -> str:
|
||||
"""Retrieve documents considering conversation history."""
|
||||
query = input_dict["query"]
|
||||
conversation = input_dict["conversation"]
|
||||
|
||||
# You could enhance this to consider historical context in retrieval
|
||||
relevant_docs = self.search_documents(query, conversation.workspace)
|
||||
if not relevant_docs:
|
||||
print("didn't find any relevant docs")
|
||||
return relevant_docs
|
||||
else:
|
||||
return relevant_docs
|
||||
|
||||
|
||||
def search_documents(self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4) -> List[Document]:
|
||||
"""Search relevant documents from the vector store."""
|
||||
filter_dict = {}
|
||||
if workspace:
|
||||
filter_dict["workspace_id"] = workspace.id
|
||||
print(f"search_kwargs: {search_kwargs}")
|
||||
retriever = self.vector_store.as_retriever(
|
||||
search_type="similarity",
|
||||
search_kwargs={
|
||||
"k": k,
|
||||
"filter": filter_dict if filter_dict else None
|
||||
}
|
||||
)
|
||||
return retriever.get_relevant_documents(query)
|
||||
|
||||
def generate_response(self, conversation: Conversation, query: str, **kwargs) -> Generator[str, None, None]:
|
||||
"""Generate response with streaming support."""
|
||||
chain_input = {
|
||||
"query": query,
|
||||
"conversation": conversation
|
||||
}
|
||||
|
||||
for chunk in self.rag_chain.stream(chain_input):
|
||||
yield chunk
|
||||
|
||||
|
||||
class AsyncRAGService(RAGService):
|
||||
"""Asynchronous RAG service implementation."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._setup_chain()
|
||||
|
||||
def _setup_chain(self):
|
||||
"""Setup the RAG chain."""
|
||||
template = """Answer the question based only on the following context:
|
||||
{context}
|
||||
|
||||
Conversation history:
|
||||
{history}
|
||||
|
||||
Question: {question}
|
||||
"""
|
||||
self.prompt = ChatPromptTemplate.from_template(template)
|
||||
|
||||
self.rag_chain = (
|
||||
{
|
||||
"context": self._retriever_with_history,
|
||||
"history": lambda x: self._format_history(x["conversation"]),
|
||||
"question": lambda x: x["query"]
|
||||
}
|
||||
| self.prompt
|
||||
| self.llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
async def _format_history(self, conversation: Conversation) -> str:
|
||||
"""Format conversation history for the prompt."""
|
||||
prompts = await Prompt.objects.filter(conversation=conversation).order_by('created_at').alist()
|
||||
print(f"prompts that we are seeding with are: {prompts}")
|
||||
return "\n".join(
|
||||
f"{'User' if prompt.is_user else 'AI'}: {prompt.text}"
|
||||
for prompt in prompts
|
||||
)
|
||||
|
||||
async def _retriever_with_history(self, input_dict: Dict[str, Any]) -> str:
|
||||
"""Retrieve documents considering conversation history."""
|
||||
print(f"Retrieving history with input: {input_dict}")
|
||||
query = input_dict["query"]
|
||||
conversation = input_dict["conversation"]
|
||||
workspace = input_dict["workspace"]
|
||||
|
||||
# You could enhance this to consider historical context in retrieval
|
||||
docs= await self.search_documents(query, workspace)
|
||||
|
||||
if not docs:
|
||||
print("Didn't find any relevant docs")
|
||||
|
||||
print("\n\n".join(doc.page_content for doc in docs))
|
||||
return "\n\n".join(doc.page_content for doc in docs)
|
||||
|
||||
|
||||
async def search_documents(self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4) -> List[Document]:
|
||||
"""Search relevant documents from the vector store."""
|
||||
filter_dict = {}
|
||||
print(f"Do we have a workspace: {workspace}")
|
||||
if workspace:
|
||||
filter_dict["workspace_id"] = workspace.id
|
||||
search_kwargs={
|
||||
"k": k,
|
||||
"filter": filter_dict if filter_dict else None
|
||||
}
|
||||
print(f"search_kwargs: {search_kwargs}")
|
||||
|
||||
retriever = self.vector_store.as_retriever(
|
||||
search_type="mmr",
|
||||
search_kwargs={
|
||||
"k": k,
|
||||
"filter": filter_dict if filter_dict else None
|
||||
}
|
||||
)
|
||||
return await retriever.aget_relevant_documents(query)
|
||||
|
||||
async def generate_response(self, conversation: Conversation, query: str, workspace: DocumentWorkspace, **kwargs) -> AsyncGenerator[str, None]:
|
||||
"""Generate response with streaming support."""
|
||||
chain_input = {
|
||||
"query": query,
|
||||
"conversation": conversation,
|
||||
"workspace": workspace,
|
||||
}
|
||||
|
||||
async for chunk in self.rag_chain.astream(chain_input):
|
||||
yield chunk
|
||||
@@ -0,0 +1,219 @@
|
||||
import os
|
||||
from unittest import TestCase, mock
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from django.test import TestCase as DjangoTestCase
|
||||
|
||||
from chat_backend.services.rag_services import RAGService, SyncRAGService, AsyncRAGService
|
||||
from chat_backend.models import Conversation, Prompt, DocumentWorkspace, Document
|
||||
|
||||
class TestRAGService(TestCase):
|
||||
def setUp(self):
|
||||
self.rag_service = RAGService()
|
||||
self.rag_service.vector_store = MagicMock()
|
||||
self.rag_service.embedding_model = MagicMock()
|
||||
self.rag_service.text_splitter = MagicMock()
|
||||
|
||||
def test_initialize_vector_store(self):
|
||||
with patch('os.path.exists', return_value=False), \
|
||||
patch('os.makedirs') as mock_makedirs, \
|
||||
patch('langchain_community.vectorstores.Chroma') as mock_chroma:
|
||||
|
||||
# Reset the vector store to test initialization
|
||||
self.rag_service.vector_store = None
|
||||
result = self.rag_service._initialize_vector_store()
|
||||
|
||||
mock_makedirs.assert_called_once_with("chroma_db")
|
||||
mock_chroma.assert_called_once_with(
|
||||
embedding_function=self.rag_service.embedding_model,
|
||||
persist_directory="chroma_db"
|
||||
)
|
||||
self.assertIsNotNone(result)
|
||||
|
||||
def test_prepare_documents(self):
|
||||
mock_doc1 = MagicMock(spec=Document)
|
||||
mock_doc1.content = "Test content"
|
||||
mock_doc1.source = "test_source"
|
||||
mock_doc1.workspace = MagicMock()
|
||||
mock_doc1.workspace.id = 1
|
||||
mock_doc1.id = 1
|
||||
|
||||
self.rag_service.text_splitter.split_text.return_value = ["chunk1", "chunk2"]
|
||||
|
||||
result = self.rag_service._prepare_documents([mock_doc1])
|
||||
|
||||
self.assertEqual(len(result), 2)
|
||||
self.rag_service.text_splitter.split_text.assert_called_once_with("Test content")
|
||||
self.assertEqual(result[0].page_content, "chunk1")
|
||||
self.assertEqual(result[0].metadata["source"], "test_source")
|
||||
|
||||
def test_ingest_documents(self):
|
||||
mock_workspace = MagicMock()
|
||||
mock_document = MagicMock()
|
||||
mock_documents = [mock_document]
|
||||
|
||||
with patch('services.rag_services.Document.objects.filter', return_value=mock_documents):
|
||||
self.rag_service._prepare_documents = MagicMock(return_value=["processed_doc"])
|
||||
|
||||
self.rag_service.ingest_documents(mock_workspace)
|
||||
|
||||
self.rag_service.vector_store.add_documents.assert_called_once_with(["processed_doc"])
|
||||
self.rag_service.vector_store.persist.assert_called_once()
|
||||
|
||||
|
||||
class TestSyncRAGService(DjangoTestCase):
|
||||
def setUp(self):
|
||||
self.sync_service = SyncRAGService()
|
||||
self.sync_service.vector_store = MagicMock()
|
||||
self.sync_service.llm = MagicMock()
|
||||
self.sync_service.rag_chain = MagicMock()
|
||||
|
||||
self.mock_conversation = MagicMock(spec=Conversation)
|
||||
self.mock_conversation.workspace = MagicMock()
|
||||
|
||||
self.mock_prompt1 = MagicMock(spec=Prompt)
|
||||
self.mock_prompt1.is_user = True
|
||||
self.mock_prompt1.text = "User question"
|
||||
self.mock_prompt1.created_at = "2023-01-01"
|
||||
|
||||
self.mock_prompt2 = MagicMock(spec=Prompt)
|
||||
self.mock_prompt2.is_user = False
|
||||
self.mock_prompt2.text = "AI response"
|
||||
self.mock_prompt2.created_at = "2023-01-02"
|
||||
|
||||
def test_format_history(self):
|
||||
with patch('services.rag_services.Prompt.objects.filter') as mock_filter:
|
||||
mock_filter.return_value.order_by.return_value = [self.mock_prompt1, self.mock_prompt2]
|
||||
|
||||
result = self.sync_service._format_history(self.mock_conversation)
|
||||
|
||||
expected = "User: User question\nAI: AI response"
|
||||
self.assertEqual(result, expected)
|
||||
mock_filter.assert_called_once_with(conversation=self.mock_conversation)
|
||||
|
||||
def test_retriever_with_history(self):
|
||||
input_dict = {
|
||||
"query": "test query",
|
||||
"conversation": self.mock_conversation
|
||||
}
|
||||
|
||||
self.sync_service.search_documents = MagicMock(return_value=["doc1", "doc2"])
|
||||
|
||||
result = self.sync_service._retriever_with_history(input_dict)
|
||||
|
||||
self.sync_service.search_documents.assert_called_once_with(
|
||||
"test query",
|
||||
self.mock_conversation.workspace
|
||||
)
|
||||
self.assertEqual(result, ["doc1", "doc2"])
|
||||
|
||||
def test_search_documents(self):
|
||||
mock_retriever = MagicMock()
|
||||
mock_retriever.get_relevant_documents.return_value = ["doc1", "doc2"]
|
||||
self.sync_service.vector_store.as_retriever.return_value = mock_retriever
|
||||
|
||||
result = self.sync_service.search_documents("test query", self.mock_conversation.workspace)
|
||||
|
||||
self.sync_service.vector_store.as_retriever.assert_called_once_with(
|
||||
search_type="similarity",
|
||||
search_kwargs={
|
||||
"k": 4,
|
||||
"filter": {"workspace_id": self.mock_conversation.workspace.id}
|
||||
}
|
||||
)
|
||||
self.assertEqual(result, ["doc1", "doc2"])
|
||||
|
||||
def test_generate_response(self):
|
||||
chain_input = {
|
||||
"query": "test query",
|
||||
"conversation": self.mock_conversation
|
||||
}
|
||||
|
||||
mock_stream = ["chunk1", "chunk2", "chunk3"]
|
||||
self.sync_service.rag_chain.stream.return_value = mock_stream
|
||||
|
||||
result = list(self.sync_service.generate_response(self.mock_conversation, "test query"))
|
||||
|
||||
self.sync_service.rag_chain.stream.assert_called_once_with(chain_input)
|
||||
self.assertEqual(result, mock_stream)
|
||||
|
||||
|
||||
class TestAsyncRAGService(DjangoTestCase):
|
||||
def setUp(self):
|
||||
self.async_service = AsyncRAGService()
|
||||
self.async_service.vector_store = MagicMock()
|
||||
self.async_service.llm = MagicMock()
|
||||
self.async_service.rag_chain = AsyncMock()
|
||||
|
||||
self.mock_conversation = MagicMock(spec=Conversation)
|
||||
self.mock_conversation.workspace = MagicMock()
|
||||
|
||||
self.mock_prompt1 = MagicMock(spec=Prompt)
|
||||
self.mock_prompt1.is_user = True
|
||||
self.mock_prompt1.text = "User question"
|
||||
self.mock_prompt1.created_at = "2023-01-01"
|
||||
|
||||
self.mock_prompt2 = MagicMock(spec=Prompt)
|
||||
self.mock_prompt2.is_user = False
|
||||
self.mock_prompt2.text = "AI response"
|
||||
self.mock_prompt2.created_at = "2023-01-02"
|
||||
|
||||
async def test_format_history(self):
|
||||
mock_manager = AsyncMock()
|
||||
mock_manager.order_by.return_value.alist.return_value = [self.mock_prompt1, self.mock_prompt2]
|
||||
|
||||
with patch('services.rag_services.Prompt.objects.filter', return_value=mock_manager):
|
||||
result = await self.async_service._format_history(self.mock_conversation)
|
||||
|
||||
expected = "User: User question\nAI: AI response"
|
||||
self.assertEqual(result, expected)
|
||||
mock_manager.order_by.assert_called_once_with('created_at')
|
||||
|
||||
async def test_retriever_with_history(self):
|
||||
input_dict = {
|
||||
"query": "test query",
|
||||
"conversation": self.mock_conversation
|
||||
}
|
||||
|
||||
self.async_service.search_documents = AsyncMock(return_value=["doc1", "doc2"])
|
||||
|
||||
result = await self.async_service._retriever_with_history(input_dict)
|
||||
|
||||
self.async_service.search_documents.assert_awaited_once_with(
|
||||
"test query",
|
||||
self.mock_conversation.workspace
|
||||
)
|
||||
self.assertEqual(result, ["doc1", "doc2"])
|
||||
|
||||
async def test_search_documents(self):
|
||||
mock_retriever = AsyncMock()
|
||||
mock_retriever.aget_relevant_documents.return_value = ["doc1", "doc2"]
|
||||
self.async_service.vector_store.as_retriever.return_value = mock_retriever
|
||||
|
||||
result = await self.async_service.search_documents("test query", self.mock_conversation.workspace)
|
||||
|
||||
self.async_service.vector_store.as_retriever.assert_called_once_with(
|
||||
search_type="similarity",
|
||||
search_kwargs={
|
||||
"k": 4,
|
||||
"filter": {"workspace_id": self.mock_conversation.workspace.id}
|
||||
}
|
||||
)
|
||||
self.assertEqual(result, ["doc1", "doc2"])
|
||||
|
||||
async def test_generate_response(self):
|
||||
chain_input = {
|
||||
"query": "test query",
|
||||
"conversation": self.mock_conversation
|
||||
}
|
||||
|
||||
mock_stream = ["chunk1", "chunk2", "chunk3"]
|
||||
self.async_service.rag_chain.astream.return_value = mock_stream
|
||||
|
||||
chunks = []
|
||||
async for chunk in self.async_service.generate_response(self.mock_conversation, "test query"):
|
||||
chunks.append(chunk)
|
||||
|
||||
self.async_service.rag_chain.astream.assert_awaited_once_with(chain_input)
|
||||
self.assertEqual(chunks, mock_stream)
|
||||
@@ -0,0 +1,67 @@
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_community.llms import Ollama
|
||||
from typing import Optional
|
||||
|
||||
class TitleGenerator:
|
||||
"""
|
||||
Generates short, descriptive titles for conversations based on the first prompt.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.llm = Ollama(
|
||||
model="llama3",
|
||||
temperature=0.5, # Slightly creative but not too random
|
||||
top_k=20,
|
||||
num_ctx=2048 # Shorter context needed for titles
|
||||
)
|
||||
|
||||
self.title_prompt = ChatPromptTemplate.from_messages([
|
||||
("system", """You are a conversation title generator. Create a very short (2-5 word) title based on the user's first message.
|
||||
|
||||
Rules:
|
||||
1. Keep it extremely concise
|
||||
2. Capture the main topic or intent
|
||||
3. Use title case
|
||||
4. No quotes or punctuation
|
||||
5. Never exceed 5 words
|
||||
|
||||
Examples:
|
||||
- "What's the weather today?" → "Weather Inquiry"
|
||||
- "Explain quantum computing" → "Quantum Computing Explanation"
|
||||
- "Generate an image of a dragon" → "Dragon Image Generation"
|
||||
- "Find our company's privacy policy" → "Privacy Policy Search"
|
||||
|
||||
Return ONLY the title, nothing else."""),
|
||||
("human", "{prompt}")
|
||||
])
|
||||
|
||||
self.chain = self.title_prompt | self.llm
|
||||
|
||||
async def generate_async(self, prompt: str) -> str:
|
||||
"""Generate title asynchronously"""
|
||||
try:
|
||||
response = await self.chain.ainvoke({"prompt": prompt})
|
||||
return self._clean_response(response)
|
||||
except Exception as e:
|
||||
print(f"Title generation error: {e}")
|
||||
return "Conversation"
|
||||
|
||||
def generate(self, prompt: str) -> str:
|
||||
"""Generate title synchronously"""
|
||||
try:
|
||||
response = self.chain.invoke({"prompt": prompt})
|
||||
return self._clean_response(response)
|
||||
except Exception as e:
|
||||
print(f"Title generation error: {e}")
|
||||
return "Conversation"
|
||||
|
||||
def _clean_response(self, response: str) -> str:
|
||||
"""Clean and format the LLM response"""
|
||||
# Remove any quotes or punctuation
|
||||
response = response.strip('"\'.!? \n\t')
|
||||
# Ensure title case and trim
|
||||
return response.title()[:50] # Hard limit for safety
|
||||
|
||||
|
||||
# Singleton instance
|
||||
title_generator = TitleGenerator()
|
||||
Reference in New Issue
Block a user