wehub-resource-sync
579 lines
22 KiB
Python
579 lines
22 KiB
Python
import os
|
|
import glob
|
|
import logging
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any, Optional
|
|
from dataclasses import dataclass, field
|
|
|
|
import openai
|
|
import sounddevice as sd
|
|
import soundfile as sf
|
|
import assemblyai as aai
|
|
from PyPDF2 import PdfReader
|
|
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
|
|
from crewai import Agent, Task, Crew
|
|
from crewai.tools import BaseTool
|
|
from crewai.flow.flow import Flow, start, listen
|
|
|
|
import config
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Initialize clients
|
|
openai_client = openai.OpenAI(api_key=config.OPENAI_API_KEY)
|
|
aai.settings.api_key = config.ASSEMBLYAI_API_KEY
|
|
|
|
# Global variables for caching
|
|
_milvus_connection = None
|
|
_collection = None
|
|
|
|
# File patterns
|
|
FILE_PATTERNS = ["*.pdf", "*.mp3", "*.wav", "*.m4a", "*.flac", "*.txt", "*.md"]
|
|
AUDIO_EXTENSIONS = {'.mp3', '.wav', '.m4a', '.flac'}
|
|
TEXT_EXTENSIONS = {'.txt', '.md'}
|
|
|
|
|
|
@dataclass
|
|
class DataIngestionState:
|
|
"""State for data ingestion flow"""
|
|
collection: Optional[Collection] = None
|
|
chunks: List[Dict[str, Any]] = field(default_factory=list)
|
|
processed_files: List[str] = field(default_factory=list)
|
|
discovered_files: List[str] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class QueryState:
|
|
"""State for query processing flow"""
|
|
query: str = ""
|
|
transcribed_query: str = ""
|
|
search_results: str = ""
|
|
final_response: str = ""
|
|
audio_file: Optional[str] = None
|
|
|
|
|
|
def get_milvus_connection():
|
|
"""Get or create Milvus connection"""
|
|
global _milvus_connection
|
|
if _milvus_connection is None:
|
|
try:
|
|
_milvus_connection = connections.connect(host=config.MILVUS_HOST, port=config.MILVUS_PORT)
|
|
logger.info("📡 Connected to Milvus")
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect to Milvus: {e}")
|
|
raise
|
|
return _milvus_connection
|
|
|
|
|
|
def get_collection():
|
|
"""Get or create collection"""
|
|
global _collection
|
|
if _collection is None:
|
|
get_milvus_connection()
|
|
_collection = Collection(config.COLLECTION_NAME)
|
|
_collection.load()
|
|
return _collection
|
|
|
|
|
|
def transcribe_audio_file(audio_file: str) -> str:
|
|
"""Transcribe audio file using AssemblyAI"""
|
|
transcriber = aai.Transcriber()
|
|
transcript = transcriber.transcribe(audio_file)
|
|
return transcript.text
|
|
|
|
|
|
def search_vector_database(query: str, limit: int = 5) -> List[Dict[str, Any]]:
|
|
"""Search vector database for relevant information"""
|
|
try:
|
|
collection = get_collection()
|
|
|
|
# Generate query embedding
|
|
response = openai_client.embeddings.create(model=config.EMBEDDING_MODEL, input=[query])
|
|
query_embedding = response.data[0].embedding
|
|
|
|
# Search
|
|
search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
|
|
results = collection.search(
|
|
data=[query_embedding],
|
|
anns_field="embedding",
|
|
param=search_params,
|
|
limit=limit,
|
|
output_fields=["text", "source", "content_type"]
|
|
)
|
|
|
|
# Format results
|
|
search_results = []
|
|
for hits in results:
|
|
for hit in hits:
|
|
search_results.append({
|
|
"text": hit.entity.get("text"),
|
|
"source": hit.entity.get("source"),
|
|
"content_type": hit.entity.get("content_type"),
|
|
"score": hit.score
|
|
})
|
|
|
|
return search_results
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error searching vector database: {e}")
|
|
return []
|
|
|
|
|
|
def format_search_results(search_results: List[Dict[str, Any]]) -> str:
|
|
"""Format search results into readable string"""
|
|
if not search_results:
|
|
return "No relevant documents found."
|
|
|
|
formatted_results = []
|
|
for result in search_results:
|
|
relevance = (1 - result['score']) * 100
|
|
formatted_results.append(
|
|
f"Source: {result['source']} ({result['content_type']})\n"
|
|
f"Relevance: {relevance:.1f}%\n"
|
|
f"Content: {result['text'][:200]}...\n"
|
|
f"---"
|
|
)
|
|
|
|
return "\n".join(formatted_results)
|
|
|
|
|
|
class SearchKnowledgeBaseTool(BaseTool):
|
|
"""Tool for searching the knowledge base"""
|
|
|
|
name: str = "search_knowledge_base"
|
|
description: str = "Search the multimodal knowledge base for relevant information"
|
|
|
|
def _run(self, query: str) -> str:
|
|
"""Search the knowledge base and return formatted results"""
|
|
logger.info(f"Searching knowledge base for: {query}")
|
|
search_results = search_vector_database(query)
|
|
return format_search_results(search_results)
|
|
|
|
|
|
class DataIngestionFlow(Flow):
|
|
"""CrewAI Flow for data ingestion and vector database setup"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
@start()
|
|
def discover_files(self) -> DataIngestionState:
|
|
"""Discover all files in the data directory"""
|
|
logger.info("🔍 Discovering files in data directory...")
|
|
data_dir = Path(config.DATA_DIR)
|
|
|
|
if not data_dir.exists():
|
|
raise FileNotFoundError(f"Data directory not found: {data_dir}")
|
|
|
|
discovered_files = []
|
|
for pattern in FILE_PATTERNS:
|
|
files = glob.glob(str(data_dir / pattern))
|
|
discovered_files.extend(files)
|
|
|
|
discovered_files = sorted(list(set(discovered_files)))
|
|
logger.info(f"📁 Discovered {len(discovered_files)} files")
|
|
return DataIngestionState(discovered_files=discovered_files)
|
|
|
|
@listen(lambda state: state.discovered_files is not None)
|
|
def setup_vector_database(self, state: DataIngestionState) -> DataIngestionState:
|
|
"""Initialize Milvus connection and collection"""
|
|
logger.info("🔧 Setting up vector database...")
|
|
|
|
get_milvus_connection()
|
|
|
|
if not utility.has_collection(config.COLLECTION_NAME):
|
|
fields = [
|
|
FieldSchema("id", DataType.INT64, is_primary=True, auto_id=True),
|
|
FieldSchema("text", DataType.VARCHAR, max_length=65535),
|
|
FieldSchema("source", DataType.VARCHAR, max_length=255),
|
|
FieldSchema("content_type", DataType.VARCHAR, max_length=50),
|
|
FieldSchema("embedding", DataType.FLOAT_VECTOR, dim=config.EMBEDDING_DIM)
|
|
]
|
|
schema = CollectionSchema(fields, "Multimodal RAG collection")
|
|
collection = Collection(config.COLLECTION_NAME, schema)
|
|
collection.create_index("embedding", {"index_type": "IVF_FLAT", "metric_type": "L2", "params": {"nlist": 1024}})
|
|
logger.info("✅ Collection created and indexed")
|
|
else:
|
|
logger.info("✅ Collection already exists")
|
|
collection = Collection(config.COLLECTION_NAME)
|
|
|
|
collection.load()
|
|
logger.info("✅ Vector database setup completed")
|
|
return DataIngestionState(discovered_files=state.discovered_files, collection=collection)
|
|
|
|
@listen(lambda state: state.collection is not None)
|
|
def process_multimodal_data(self, state: DataIngestionState) -> DataIngestionState:
|
|
"""Process discovered files from the data directory"""
|
|
logger.info(f"📄 Processing {len(state.discovered_files)} discovered files...")
|
|
|
|
chunks = []
|
|
processed_files = []
|
|
|
|
for file_path in state.discovered_files:
|
|
file_path = Path(file_path)
|
|
filename = file_path.name
|
|
logger.info(f"🔄 Processing: {filename}")
|
|
|
|
try:
|
|
# Process different file types
|
|
if file_path.suffix.lower() == '.pdf':
|
|
with open(file_path, 'rb') as f:
|
|
reader = PdfReader(f)
|
|
text = "\n".join(page.extract_text() for page in reader.pages)
|
|
content_type = "pdf"
|
|
elif file_path.suffix.lower() in AUDIO_EXTENSIONS:
|
|
text = transcribe_audio_file(str(file_path))
|
|
content_type = "audio"
|
|
elif file_path.suffix.lower() in TEXT_EXTENSIONS:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
text = f.read()
|
|
content_type = "text"
|
|
else:
|
|
logger.warning(f"⚠️ Skipping unsupported file: {filename}")
|
|
continue
|
|
|
|
# Create chunks
|
|
chunk_size = 1000
|
|
for i in range(0, len(text), chunk_size):
|
|
chunk = text[i:i + chunk_size]
|
|
if chunk.strip():
|
|
chunks.append({
|
|
"text": chunk,
|
|
"source": filename,
|
|
"content_type": content_type
|
|
})
|
|
|
|
processed_files.append(filename)
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Error processing {filename}: {e}")
|
|
continue
|
|
|
|
if not chunks:
|
|
logger.warning("No content extracted from files")
|
|
return DataIngestionState(
|
|
discovered_files=state.discovered_files,
|
|
collection=state.collection,
|
|
chunks=[],
|
|
processed_files=[]
|
|
)
|
|
|
|
logger.info(f"✅ Processed {len(chunks)} chunks from {len(processed_files)} files")
|
|
return DataIngestionState(
|
|
discovered_files=state.discovered_files,
|
|
collection=state.collection,
|
|
chunks=chunks,
|
|
processed_files=processed_files
|
|
)
|
|
|
|
@listen(lambda state: len(state.chunks) > 0)
|
|
def generate_embeddings_flow(self, state: DataIngestionState) -> DataIngestionState:
|
|
"""Generate embeddings for processed chunks"""
|
|
logger.info(f"🧠 Generating embeddings for {len(state.chunks)} chunks...")
|
|
|
|
texts = [chunk["text"] for chunk in state.chunks]
|
|
batch_size = 100
|
|
all_embeddings = []
|
|
|
|
for i in range(0, len(texts), batch_size):
|
|
batch_texts = texts[i:i+batch_size]
|
|
response = openai_client.embeddings.create(model=config.EMBEDDING_MODEL, input=batch_texts)
|
|
batch_embeddings = [data.embedding for data in response.data]
|
|
all_embeddings.extend(batch_embeddings)
|
|
|
|
# Assign embeddings to chunks
|
|
updated_chunks = []
|
|
for chunk, embedding in zip(state.chunks, all_embeddings):
|
|
chunk_copy = chunk.copy()
|
|
chunk_copy["embedding"] = embedding
|
|
updated_chunks.append(chunk_copy)
|
|
|
|
logger.info("✅ Embeddings generation completed")
|
|
return DataIngestionState(
|
|
discovered_files=state.discovered_files,
|
|
collection=state.collection,
|
|
chunks=updated_chunks,
|
|
processed_files=state.processed_files
|
|
)
|
|
|
|
@listen(lambda state: all(chunk.get("embedding") is not None for chunk in state.chunks))
|
|
def store_in_vector_database(self, state: DataIngestionState) -> DataIngestionState:
|
|
"""Insert processed chunks into Milvus"""
|
|
logger.info(f"💾 Inserting {len(state.chunks)} chunks into vector database...")
|
|
|
|
data = [
|
|
[chunk["text"] for chunk in state.chunks],
|
|
[chunk["source"] for chunk in state.chunks],
|
|
[chunk["content_type"] for chunk in state.chunks],
|
|
[chunk["embedding"] for chunk in state.chunks]
|
|
]
|
|
|
|
state.collection.insert(data)
|
|
state.collection.flush()
|
|
logger.info("✅ Data insertion completed")
|
|
return state
|
|
|
|
|
|
class MultimodalRAGFlow(Flow):
|
|
"""CrewAI Flow for query processing and response generation"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
@start()
|
|
def transcribe_audio_if_needed(self, query: str, audio_file: Optional[str] = None) -> QueryState:
|
|
"""Transcribe audio if query is audio-based"""
|
|
if audio_file:
|
|
logger.info("🎤 Transcribing audio file...")
|
|
transcribed_query = transcribe_audio_file(audio_file)
|
|
logger.info(f"✅ Transcribed: {transcribed_query}")
|
|
else:
|
|
transcribed_query = query
|
|
logger.info("📝 Using text query directly")
|
|
|
|
return QueryState(query=query, audio_file=audio_file, transcribed_query=transcribed_query)
|
|
|
|
@listen(lambda state: state.transcribed_query is not None)
|
|
def search_knowledge_base(self, state: QueryState) -> QueryState:
|
|
"""Search the vector database for relevant information"""
|
|
logger.info("🔍 Searching knowledge base...")
|
|
|
|
search_results = search_vector_database(state.transcribed_query)
|
|
formatted_results = format_search_results(search_results)
|
|
logger.info("✅ Search completed")
|
|
|
|
return QueryState(
|
|
query=state.query,
|
|
audio_file=state.audio_file,
|
|
transcribed_query=state.transcribed_query,
|
|
search_results=formatted_results
|
|
)
|
|
|
|
@listen(lambda state: state.search_results is not None)
|
|
def generate_response(self, state: QueryState) -> QueryState:
|
|
"""Generate final response using CrewAI agents"""
|
|
logger.info("🤖 Generating response...")
|
|
|
|
try:
|
|
# Create agents
|
|
research_agent = Agent(
|
|
role="Information Retrieval Specialist",
|
|
goal="Find the most relevant information from the knowledge base to answer user queries",
|
|
backstory="You are an expert at analyzing queries and searching through multimodal knowledge bases to find the most relevant information.",
|
|
tools=[SearchKnowledgeBaseTool()],
|
|
verbose=True,
|
|
allow_delegation=False
|
|
)
|
|
|
|
response_agent = Agent(
|
|
role="Response Generator",
|
|
goal="Generate comprehensive, accurate, and helpful responses based on retrieved information",
|
|
backstory="You are an expert at synthesizing information from multiple sources and creating clear, informative responses.",
|
|
verbose=True,
|
|
allow_delegation=False
|
|
)
|
|
|
|
# Create tasks
|
|
research_task = Task(
|
|
description=f"Search for information relevant to: '{state.transcribed_query}'. Use the search_knowledge_base tool to find the most relevant context.",
|
|
agent=research_agent,
|
|
expected_output="Detailed information from the knowledge base with proper citations"
|
|
)
|
|
|
|
response_task = Task(
|
|
description=f"Based on the research findings, generate a comprehensive response to: '{state.transcribed_query}'.",
|
|
agent=response_agent,
|
|
expected_output="A well-structured, comprehensive response with proper citations"
|
|
)
|
|
|
|
# Create crew and execute
|
|
crew = Crew(
|
|
agents=[research_agent, response_agent],
|
|
tasks=[research_task, response_task],
|
|
verbose=True,
|
|
memory=False
|
|
)
|
|
|
|
result = crew.kickoff()
|
|
logger.info("✅ Response generated")
|
|
|
|
return QueryState(
|
|
query=state.query,
|
|
audio_file=state.audio_file,
|
|
transcribed_query=state.transcribed_query,
|
|
search_results=state.search_results,
|
|
final_response=result
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error generating response: {e}")
|
|
# Fallback response
|
|
fallback_response = f"Based on the search results:\n{state.search_results}\n\nQuery: {state.transcribed_query}"
|
|
return QueryState(
|
|
query=state.query,
|
|
audio_file=state.audio_file,
|
|
transcribed_query=state.transcribed_query,
|
|
search_results=state.search_results,
|
|
final_response=fallback_response
|
|
)
|
|
|
|
|
|
def record_audio(duration: int = 10, sample_rate: int = 16000) -> str:
|
|
"""Record audio from microphone and save to temporary file"""
|
|
print(f"🎤 Recording for {duration} seconds... Speak now!")
|
|
print("Press Ctrl+C to stop early")
|
|
|
|
try:
|
|
audio_data = sd.rec(int(duration * sample_rate), samplerate=sample_rate, channels=1, dtype='float64')
|
|
sd.wait()
|
|
|
|
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
|
|
sf.write(temp_file.name, audio_data, sample_rate)
|
|
|
|
print("✅ Recording completed!")
|
|
return temp_file.name
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n⏹️ Recording stopped by user")
|
|
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
|
|
sf.write(temp_file.name, audio_data, sample_rate)
|
|
return temp_file.name
|
|
except Exception as e:
|
|
print(f"❌ Recording failed: {e}")
|
|
raise
|
|
|
|
|
|
def process_query(query: str, audio_file: Optional[str] = None) -> str:
|
|
"""Process a query through the RAG flow"""
|
|
try:
|
|
rag_flow = MultimodalRAGFlow()
|
|
|
|
# Start the flow and execute all steps
|
|
initial_state = rag_flow.transcribe_audio_if_needed(query, audio_file)
|
|
|
|
# Manually execute the remaining steps since Flow isn't auto-executing
|
|
search_state = rag_flow.search_knowledge_base(initial_state)
|
|
final_state = rag_flow.generate_response(search_state)
|
|
|
|
return final_state.final_response
|
|
except Exception as e:
|
|
logger.error(f"Error processing query: {e}")
|
|
return f"Sorry, I encountered an error while processing your query: {str(e)}"
|
|
|
|
|
|
def check_system_status():
|
|
"""Check if the system is ready"""
|
|
try:
|
|
# Check Milvus connection
|
|
get_milvus_connection()
|
|
|
|
# Check if collection exists and has data
|
|
if utility.has_collection(config.COLLECTION_NAME):
|
|
collection = Collection(config.COLLECTION_NAME)
|
|
collection.load()
|
|
# Check if collection has data
|
|
stats = collection.get_statistics()
|
|
if stats.get("row_count", 0) > 0:
|
|
logger.info("✅ System ready with data")
|
|
return True
|
|
else:
|
|
logger.info("⚠️ Collection exists but is empty")
|
|
return False
|
|
else:
|
|
logger.info("⚠️ Collection does not exist")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"❌ System check failed: {e}")
|
|
return False
|
|
|
|
|
|
def main():
|
|
"""Main application entry point"""
|
|
print("\n🤖 Welcome to Multimodal Agentic RAG System!")
|
|
print("=" * 50)
|
|
|
|
# Validate API keys
|
|
if not config.ASSEMBLYAI_API_KEY or not config.OPENAI_API_KEY:
|
|
print("\n❌ Missing API keys!")
|
|
print("Please create a .env file with:")
|
|
print("ASSEMBLYAI_API_KEY=your_key_here")
|
|
print("OPENAI_API_KEY=your_key_here")
|
|
return
|
|
|
|
# Check system status and setup if needed
|
|
print("🔍 Checking system status...")
|
|
try:
|
|
if check_system_status():
|
|
print("✅ System ready!")
|
|
else:
|
|
print("\n⚠️ System not set up yet. Let's set it up first!")
|
|
print("📡 Connecting to Milvus...")
|
|
|
|
# Use the DataIngestionFlow to set up the system
|
|
ingestion_flow = DataIngestionFlow()
|
|
initial_state = ingestion_flow.discover_files()
|
|
|
|
# Manually execute the flow steps
|
|
setup_state = ingestion_flow.setup_vector_database(initial_state)
|
|
process_state = ingestion_flow.process_multimodal_data(setup_state)
|
|
|
|
if len(process_state.chunks) > 0:
|
|
embed_state = ingestion_flow.generate_embeddings_flow(process_state)
|
|
final_state = ingestion_flow.store_in_vector_database(embed_state)
|
|
print("✅ Setup completed!")
|
|
else:
|
|
print("⚠️ No data to process")
|
|
|
|
except Exception as e:
|
|
print(f"\n⚠️ Error checking system status: {e}")
|
|
print("Make sure Milvus is running: docker-compose up -d")
|
|
return
|
|
|
|
# Main interaction loop
|
|
while True:
|
|
print("\nWhat would you like to do?")
|
|
print("1. 💬 Ask a question (text)")
|
|
print("2. 🎤 Record and ask a question")
|
|
print("3. 🚪 Exit")
|
|
|
|
choice = input("\nEnter your choice (1-3): ").strip()
|
|
|
|
if choice == "1":
|
|
query = input("\n💬 Enter your question: ").strip()
|
|
if query:
|
|
print("\n🤔 Processing...")
|
|
response = process_query(query)
|
|
print(f"\n🤖 Response:\n{response}")
|
|
|
|
elif choice == "2":
|
|
duration = input("\n🎤 Recording duration in seconds (default 10): ").strip()
|
|
duration = int(duration) if duration.isdigit() else 10
|
|
|
|
audio_file = record_audio(duration)
|
|
if audio_file:
|
|
print("\n🤔 Processing...")
|
|
response = process_query("", audio_file)
|
|
print(f"\n🤖 Response:\n{response}")
|
|
|
|
try:
|
|
os.unlink(audio_file)
|
|
except:
|
|
pass
|
|
|
|
elif choice == "3":
|
|
print("\n👋 Goodbye!")
|
|
break
|
|
|
|
else:
|
|
print("❌ Invalid choice. Please enter 1-3.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except Exception as e:
|
|
print(f"❌ Fatal error: {e}")
|
|
import traceback
|
|
traceback.print_exc() |