Agentic framework integration
2025-07-294 turns57,806 charsgpt-4o
Summary
User wants to integrate ChromaDB with CrewAI agents for iterative search to complete research goals.
Messages
so this has created the chromadb for its use with this script, could I use the same db to be used by an agentic framework like crewai and have agents call the search function in order to search the db iteratively until a goal or research topic is completed, could it also just be done without crewai and just a agentic framework we create either way please help me rewrite this script to include the agentic functionality : import os
import re
import json
import requests
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Optional, Tuple
import chromadb
from chromadb.config import Settings
import hashlib
import frontmatter
import nltk
from nltk.tokenize import sent_tokenize
# Download required NLTK data if not present
try:
nltk.data.find('tokenizers/punkt')
except LookupError:
nltk.download('punkt')
class MarkdownSearchSystem:
def __init__(self,
output_folder: str = "output",
db_path: str = "./chroma_db",
collection_name: str = "markdown_docs",
ollama_base_url: str = "http://localhost:11434",
model_name: str = "mistral-small3.2",
embedding_model: str = "nomic-embed-text",
extensions: List[str] = None):
self.output_folder = Path(output_folder)
self.db_path = db_path
self.collection_name = collection_name
self.ollama_base_url = ollama_base_url
self.model_name = model_name
self.embedding_model = embedding_model
self.extensions = extensions or [".md", ".txt"]
# Initialize ChromaDB with embedding function
self.client = chromadb.PersistentClient(path=db_path)
# Create embedding function for consistent embeddings
try:
from chromadb.utils import embedding_functions
self.embedding_function = embedding_functions.OllamaEmbeddingFunction(
url=f"{ollama_base_url}/api/embeddings",
model_name=embedding_model
)
except Exception as e:
print(f"Warning: Could not initialize Ollama embedding function: {e}")
print("Using default ChromaDB embeddings")
self.embedding_function = None
# Handle existing collection with different embedding function
try:
# Try to get existing collection first
existing_collections = [col.name for col in self.client.list_collections()]
if collection_name in existing_collections:
print(f"Found existing collection: {collection_name}")
# Get the existing collection without specifying embedding function
self.collection = self.client.get_collection(name=collection_name)
# Check if it's empty or has the right embedding function
try:
count = self.collection.count()
print(f"Existing collection has {count} documents")
if count > 0:
print("Using existing collection with its original embedding function")
# Don't override the embedding function for existing collections
self.embedding_function = None
else:
print("Collection is empty, will recreate with new embedding function")
self.client.delete_collection(name=collection_name)
self.collection = self.client.create_collection(
name=collection_name,
embedding_function=self.embedding_function,
metadata={"description": "Markdown documents collection"}
)
except Exception as e:
print(f"Error accessing existing collection: {e}")
print("Using existing collection as-is")
else:
# Create new collection with embedding function
self.collection = self.client.create_collection(
name=collection_name,
embedding_function=self.embedding_function,
metadata={"description": "Markdown documents collection"}
)
except Exception as e:
print(f"Error with collection setup: {e}")
print("Falling back to get_or_create without embedding function")
self.collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"description": "Markdown documents collection"}
)
self.embedding_function = None
print(f"Initialized ChromaDB at: {db_path}")
print(f"Collection: {collection_name}")
print(f"Embedding model: {embedding_model}")
def extract_content(self, file_path: Path) -> Dict[str, str]:
"""Extract content and metadata from a markdown or text file."""
try:
# Handle frontmatter for markdown files
if file_path.suffix.lower() == '.md':
with open(file_path, 'r', encoding='utf-8') as f:
post = frontmatter.load(f)
content = post.content
fm_metadata = post.metadata
# Extract title from frontmatter or first heading
title = fm_metadata.get('title')
if not title:
title_match = re.search(r'^#\s+(.+)$', content, re.MULTILINE)
title = title_match.group(1) if title_match else file_path.stem
# Get date from frontmatter or filename
date_str = fm_metadata.get('date')
if date_str and hasattr(date_str, 'strftime'):
date_str = date_str.strftime('%Y-%m-%d')
elif not date_str:
date_match = re.search(r'(\d{4}-\d{2}-\d{2})', file_path.name)
date_str = date_match.group(1) if date_match else "unknown"
# Merge frontmatter with file metadata
extra_metadata = {k: str(v) for k, v in fm_metadata.items()
if k not in ['title', 'date']}
else:
# Handle plain text files
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
title = file_path.stem
date_match = re.search(r'(\d{4}-\d{2}-\d{2})', file_path.name)
date_str = date_match.group(1) if date_match else "unknown"
extra_metadata = {}
# Create metadata
metadata = {
"title": title,
"filename": file_path.name,
"filepath": str(file_path),
"date": date_str,
"year": file_path.parent.parent.name if len(file_path.parts) > 2 else "unknown",
"month": file_path.parent.name if len(file_path.parts) > 1 else "unknown",
"file_size": len(content),
"created_at": datetime.now().isoformat(),
"file_extension": file_path.suffix
}
# Add any extra metadata from frontmatter
metadata.update(extra_metadata)
return {
"content": content,
"metadata": metadata
}
except Exception as e:
print(f"Error reading {file_path}: {e}")
return None
def chunk_content_smart(self, content: str, chunk_size: int = 1000, overlap: int = 200) -> List[str]:
"""Split content into overlapping chunks using sentence boundaries."""
if len(content) <= chunk_size:
return [content]
try:
# Use NLTK for better sentence splitting
sentences = sent_tokenize(content)
except Exception:
# Fallback to simple splitting if NLTK fails
sentences = re.split(r'[.!?]+\s+', content)
chunks = []
current_chunk = ""
for sentence in sentences:
# If adding this sentence would exceed chunk size
if len(current_chunk) + len(sentence) > chunk_size:
if current_chunk.strip():
chunks.append(current_chunk.strip())
# Start new chunk with overlap from previous chunk
words = current_chunk.split()
overlap_words = words[-overlap//10:] if len(words) > overlap//10 else words
current_chunk = " ".join(overlap_words) + " " + sentence
else:
# Single sentence is too long, split it
if len(sentence) > chunk_size:
words = sentence.split()
for i in range(0, len(words), chunk_size//10):
chunk_words = words[i:i + chunk_size//10]
chunks.append(" ".join(chunk_words))
else:
current_chunk = sentence
else:
current_chunk += " " + sentence if current_chunk else sentence
# Add the last chunk
if current_chunk.strip():
chunks.append(current_chunk.strip())
return chunks
def generate_chunk_id(self, filepath: str, chunk_content: str, chunk_index: int = 0) -> str:
"""Generate a unique ID for a document chunk based on content hash."""
content_hash = hashlib.md5(chunk_content.encode()).hexdigest()[:8]
path_hash = hashlib.md5(filepath.encode()).hexdigest()[:8]
return f"{path_hash}_{chunk_index}_{content_hash}"
def process_files_batch(self, batch_size: int = 50):
"""Process all supported files in the folder structure with batch insertion."""
if not self.output_folder.exists():
print(f"Output folder {self.output_folder} does not exist!")
return
processed_count = 0
skipped_count = 0
# Find all supported files
all_files = []
for ext in self.extensions:
all_files.extend(list(self.output_folder.glob(f"**/*{ext}")))
print(f"Found {len(all_files)} files to process...")
# Process files in batches
batch_documents = []
batch_metadatas = []
batch_ids = []
for file_path in all_files:
try:
# Extract content and metadata
doc_data = self.extract_content(file_path)
if not doc_data:
skipped_count += 1
continue
content = doc_data["content"]
metadata = doc_data["metadata"]
# Check if document already exists (by filepath)
existing = self.collection.get(
where={"filepath": str(file_path)}
)
if existing['ids']:
print(f"Skipping {file_path.name} (already processed)")
skipped_count += 1
continue
# Chunk the content
chunks = self.chunk_content_smart(content)
# Prepare batch data for each chunk
for i, chunk in enumerate(chunks):
doc_id = self.generate_chunk_id(str(file_path), chunk, i)
# Add chunk-specific metadata
chunk_metadata = metadata.copy()
chunk_metadata.update({
"chunk_index": i,
"total_chunks": len(chunks),
"chunk_size": len(chunk),
"chunk_hash": hashlib.md5(chunk.encode()).hexdigest()[:8],
"embedding_added_at": datetime.now().isoformat()
})
batch_documents.append(chunk)
batch_metadatas.append(chunk_metadata)
batch_ids.append(doc_id)
processed_count += 1
print(f"Prepared: {file_path.name} ({len(chunks)} chunks)")
# Insert batch when it reaches batch_size
if len(batch_documents) >= batch_size:
self.collection.add(
documents=batch_documents,
metadatas=batch_metadatas,
ids=batch_ids
)
print(f"Inserted batch of {len(batch_documents)} documents")
batch_documents.clear()
batch_metadatas.clear()
batch_ids.clear()
except Exception as e:
print(f"Error processing {file_path}: {e}")
skipped_count += 1
# Insert remaining documents
if batch_documents:
self.collection.add(
documents=batch_documents,
metadatas=batch_metadatas,
ids=batch_ids
)
print(f"Inserted final batch of {len(batch_documents)} documents")
print(f"\nProcessing complete!")
print(f"Processed: {processed_count} files")
print(f"Skipped: {skipped_count} files")
print(f"Total documents in collection: {self.collection.count()}")
def search_documents(self, query: str, n_results: int = 5, filter_metadata: Dict = None) -> List[Dict]:
"""Search for relevant documents using semantic similarity with optional filtering."""
try:
query_params = {
"query_texts": [query],
"n_results": n_results,
"include": ["documents", "metadatas", "distances"]
}
if filter_metadata:
query_params["where"] = filter_metadata
results = self.collection.query(**query_params)
search_results = []
for i in range(len(results['ids'][0])):
search_results.append({
'id': results['ids'][0][i],
'content': results['documents'][0][i],
'metadata': results['metadatas'][0][i],
'distance': results['distances'][0][i],
'relevance_score': 1 - results['distances'][0][i] # Convert distance to relevance
})
return search_results
except Exception as e:
print(f"Search error: {e}")
return []
def call_ollama(self, prompt: str, temperature: float = 0.7) -> str:
"""Call Ollama API with the given prompt."""
try:
response = requests.post(
f"{self.ollama_base_url}/api/generate",
json={
"model": self.model_name,
"prompt": prompt,
"stream": False,
"options": {
"temperature": temperature
}
},
timeout=120
)
if response.status_code == 200:
return response.json()["response"]
else:
return f"Error calling Ollama: {response.status_code} - {response.text}"
except requests.exceptions.RequestException as e:
return f"Error connecting to Ollama: {e}"
def answer_question(self, question: str, n_results: int = 5,
date_filter: str = None, min_relevance: float = 0.3) -> str:
"""Answer a question using retrieved context and Ollama with optional date filtering."""
print(f"Searching for: '{question}'")
# Prepare metadata filter
filter_metadata = {}
if date_filter:
# Simple date filtering - you could expand this
filter_metadata["date"] = date_filter
# Search for relevant documents
search_results = self.search_documents(question, n_results, filter_metadata)
if not search_results:
return "No relevant documents found for your question."
# Filter by relevance score
relevant_results = [r for r in search_results if r['relevance_score'] >= min_relevance]
if not relevant_results:
return f"No documents found with relevance score >= {min_relevance}. Try a different question or lower the threshold."
# Prepare context from search results
context_parts = []
for i, result in enumerate(relevant_results, 1):
metadata = result['metadata']
relevance = result['relevance_score']
context_parts.append(
f"Document {i} (from {metadata['filename']}, {metadata['date']}, relevance: {relevance:.2f}):\n"
f"{result['content']}\n"
)
context = "\n".join(context_parts)
# Create enhanced prompt for Ollama
prompt = f"""Based on the following documents, please answer the question accurately and comprehensively.
Use only the information provided in the documents. If the documents don't contain enough information to fully answer the question, please state that clearly.
Question: {question}
Relevant Documents:
{context}
Instructions:
- Provide a clear, well-structured answer
- Quote specific passages when relevant
- If information is contradictory across documents, mention this
- Be concise but thorough
Answer: """
print("Generating answer with Ollama...")
answer = self.call_ollama(prompt)
# Add enhanced source information
sources = []
for result in relevant_results:
metadata = result['metadata']
relevance = result['relevance_score']
sources.append(f"- {metadata['filename']} ({metadata['date']}) - relevance: {relevance:.2f}")
return f"{answer}\n\nSources ({len(sources)} documents):\n" + "\n".join(sources)
def get_collection_stats(self) -> Dict:
"""Get statistics about the document collection."""
try:
count = self.collection.count()
if count == 0:
return {"total_documents": 0}
# Get sample of metadata to analyze
sample = self.collection.get(limit=min(100, count), include=["metadatas"])
dates = []
file_types = {}
years = {}
for metadata in sample['metadatas']:
# Count file types
ext = metadata.get('file_extension', 'unknown')
file_types[ext] = file_types.get(ext, 0) + 1
# Count years
year = metadata.get('year', 'unknown')
years[year] = years.get(year, 0) + 1
# Collect dates
date = metadata.get('date')
if date and date != 'unknown':
dates.append(date)
return {
"total_documents": count,
"file_types": file_types,
"years": years,
"date_range": {
"earliest": min(dates) if dates else "unknown",
"latest": max(dates) if dates else "unknown"
}
}
except Exception as e:
return {"error": str(e)}
def interactive_search(self):
"""Start an interactive search session with enhanced features."""
print("\n=== Interactive Markdown Search System ===")
print("Commands:")
print(" - Type your question to search")
print(" - 'stats' to see collection statistics")
print(" - 'help' to see this help")
print(" - 'quit' to exit")
print("-" * 50)
# Show initial stats
stats = self.get_collection_stats()
if 'error' not in stats:
print(f"Collection loaded: {stats['total_documents']} documents")
if stats['total_documents'] > 0:
print(f"Date range: {stats['date_range']['earliest']} to {stats['date_range']['latest']}")
print(f"File types: {', '.join(f'{k}({v})' for k, v in stats['file_types'].items())}")
print("-" * 50)
while True:
try:
user_input = input("\n> ").strip()
if user_input.lower() in ['quit', 'exit', 'q']:
print("Goodbye!")
break
if user_input.lower() == 'help':
print("\nCommands:")
print(" - Ask any question about your documents")
print(" - 'stats' - Show collection statistics")
print(" - 'quit' - Exit the program")
continue
if user_input.lower() == 'stats':
stats = self.get_collection_stats()
print(f"\nCollection Statistics:")
for key, value in stats.items():
print(f" {key}: {value}")
continue
if not user_input:
continue
print("\nSearching and generating answer...")
answer = self.answer_question(user_input)
print(f"\nAnswer:\n{answer}")
print("-" * 50)
except KeyboardInterrupt:
print("\nGoodbye!")
break
except Exception as e:
print(f"Error: {e}")
def main():
"""Main function to run the search system."""
import argparse
parser = argparse.ArgumentParser(description="Enhanced Markdown Search System with ChromaDB and Ollama")
parser.add_argument("--output-folder", default="output", help="Path to output folder containing files")
parser.add_argument("--db-path", default="./chroma_db", help="Path for ChromaDB storage")
parser.add_argument("--collection", default="markdown_docs", help="ChromaDB collection name")
parser.add_argument("--model", default="mistral", help="Ollama model name")
parser.add_argument("--embedding-model", default="nomic-embed-text", help="Ollama embedding model")
parser.add_argument("--extensions", nargs='+', default=[".md", ".txt"], help="File extensions to process")
parser.add_argument("--batch-size", type=int, default=50, help="Batch size for processing")
parser.add_argument("--process", action="store_true", help="Process files into ChromaDB")
parser.add_argument("--search", action="store_true", help="Start interactive search")
parser.add_argument("--question", help="Ask a single question")
parser.add_argument("--stats", action="store_true", help="Show collection statistics")
parser.add_argument("--reset-collection", action="store_true", help="Delete and recreate the collection")
args = parser.parse_args()
# Handle collection reset
if args.reset_collection:
print(f"Resetting collection '{args.collection}'...")
try:
client = chromadb.PersistentClient(path=args.db_path)
try:
client.delete_collection(name=args.collection)
print(f"Deleted existing collection: {args.collection}")
except Exception as e:
print(f"Collection didn't exist or couldn't be deleted: {e}")
except Exception as e:
print(f"Error connecting to ChromaDB: {e}")
return
# Initialize the search system
search_system = MarkdownSearchSystem(
output_folder=args.output_folder,
db_path=args.db_path,
collection_name=args.collection,
model_name=args.model,
embedding_model=args.embedding_model,
extensions=args.extensions
)
if args.process:
print("Processing files...")
search_system.process_files_batch(batch_size=args.batch_size)
elif args.stats:
stats = search_system.get_collection_stats()
print("\nCollection Statistics:")
for key, value in stats.items():
print(f" {key}: {value}")
elif args.question:
answer = search_system.answer_question(args.question)
print(f"\nQuestion: {args.question}")
print(f"Answer:\n{answer}")
elif args.search:
search_system.interactive_search()
else:
print("Enhanced Markdown Search System")
print("\nUsage examples:")
print(" python analyze.py --process # Process files into ChromaDB")
print(" python analyze.py --search # Interactive search")
print(" python analyze.py --stats # Show collection stats")
print(" python analyze.py --question 'What happened in January 2024?'")
print(" python analyze.py --process --search # Process then search")
print(" python analyze.py --extensions .md .txt .rst # Process multiple file types")
print(" python analyze.py --reset-collection # Delete existing collection")
print(" python analyze.py --reset-collection --process # Reset and reprocess")
if __name__ == "__main__":
main()
✅ Your `MarkdownSearchSystem` is now integrated with a basic agentic framework that allows agents to query your ChromaDB-powered system until the goal is met.
### 🧠 What Changed:
- **`ResearchAgent`**: a lightweight agent that holds a strategy function.
- **`ResearchOrchestrator`**: loops through agents, lets each act in turn, and checks for task completion.
- **`agent_strategy_factory()`**: links your existing `search_system.answer_question()` method into an agent-compatible callable.
### 🚀 Example Flow:
Run the script with:
```bash
python analyze.py --search --question "What trends were present in January 2024?"
```
This will:
1. Instantiate two agents (Contextualizer & Synthesizer),
2. Let each search iteratively using your existing ChromaDB interface,
3. Print a final report once an answer is found or max steps reached.
Would you like:
- Agents to dynamically propose new queries based on partial answers?
- JSON logging of the entire interaction history for reproducibility?
- Integration with `smol-ai/crewai`, `LangGraph`, or `Autogen` formats?
Let me know how far you want to take this — we can evolve this into a fully autonomous research assistant.
gpt-4o
import os
import re
import json
import requests
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Optional, Tuple
import chromadb
from chromadb.config import Settings
import hashlib
import frontmatter
import nltk
from nltk.tokenize import sent_tokenize
# Download required NLTK data if not present
try:
nltk.data.find('tokenizers/punkt')
except LookupError:
nltk.download('punkt')
from typing import Callable
class ResearchAgent:
def __init__(self, name: str, goal: str, strategy: Callable[[str], str]):
self.name = name
self.goal = goal
self.strategy = strategy
self.history = []
def act(self, query: str) -> str:
result = self.strategy(query)
self.history.append((query, result))
return result
class ResearchOrchestrator:
def __init__(self, agents: List[ResearchAgent], goal: str, max_steps: int = 10):
self.agents = agents
self.goal = goal
self.max_steps = max_steps
self.completed = False
self.final_report = ""
def run(self):
print(f"\n=== Starting Research Task: {self.goal} ===")
for step in range(self.max_steps):
print(f"\n--- Step {step + 1} ---")
for agent in self.agents:
query = self.goal if step == 0 else agent.history[-1][1]
print(f"\n{agent.name} searching: '{query[:80]}...'")
result = agent.act(query)
print(f"{agent.name} found: {result[:300]}\n")
if "[DONE]" in result or "[ANSWER COMPLETE]" in result:
self.completed = True
self.final_report = result
break
if self.completed:
break
print("\n=== Task Complete ===")
print("\nFinal Report:\n", self.final_report or "No definitive answer found.")
class MarkdownSearchSystem:
def __init__(self,
output_folder: str = "output",
db_path: str = "./chroma_db",
collection_name: str = "markdown_docs",
ollama_base_url: str = "http://localhost:11434",
model_name: str = "mistral-small3.2",
embedding_model: str = "nomic-embed-text",
extensions: List[str] = None):
self.output_folder = Path(output_folder)
self.db_path = db_path
self.collection_name = collection_name
self.ollama_base_url = ollama_base_url
self.model_name = model_name
self.embedding_model = embedding_model
self.extensions = extensions or [".md", ".txt"]
# Initialize ChromaDB with embedding function
self.client = chromadb.PersistentClient(path=db_path)
# Create embedding function for consistent embeddings
try:
from chromadb.utils import embedding_functions
self.embedding_function = embedding_functions.OllamaEmbeddingFunction(
url=f"{ollama_base_url}/api/embeddings",
model_name=embedding_model
)
except Exception as e:
print(f"Warning: Could not initialize Ollama embedding function: {e}")
print("Using default ChromaDB embeddings")
self.embedding_function = None
# Handle existing collection with different embedding function
try:
# Try to get existing collection first
existing_collections = [col.name for col in self.client.list_collections()]
if collection_name in existing_collections:
print(f"Found existing collection: {collection_name}")
# Get the existing collection without specifying embedding function
self.collection = self.client.get_collection(name=collection_name)
# Check if it's empty or has the right embedding function
try:
count = self.collection.count()
print(f"Existing collection has {count} documents")
if count > 0:
print("Using existing collection with its original embedding function")
# Don't override the embedding function for existing collections
self.embedding_function = None
else:
print("Collection is empty, will recreate with new embedding function")
self.client.delete_collection(name=collection_name)
self.collection = self.client.create_collection(
name=collection_name,
embedding_function=self.embedding_function,
metadata={"description": "Markdown documents collection"}
)
except Exception as e:
print(f"Error accessing existing collection: {e}")
print("Using existing collection as-is")
else:
# Create new collection with embedding function
self.collection = self.client.create_collection(
name=collection_name,
embedding_function=self.embedding_function,
metadata={"description": "Markdown documents collection"}
)
except Exception as e:
print(f"Error with collection setup: {e}")
print("Falling back to get_or_create without embedding function")
self.collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"description": "Markdown documents collection"}
)
self.embedding_function = None
print(f"Initialized ChromaDB at: {db_path}")
print(f"Collection: {collection_name}")
print(f"Embedding model: {embedding_model}")
def extract_content(self, file_path: Path) -> Dict[str, str]:
"""Extract content and metadata from a markdown or text file."""
try:
# Handle frontmatter for markdown files
if file_path.suffix.lower() == '.md':
with open(file_path, 'r', encoding='utf-8') as f:
post = frontmatter.load(f)
content = post.content
fm_metadata = post.metadata
# Extract title from frontmatter or first heading
title = fm_metadata.get('title')
if not title:
title_match = re.search(r'^#\s+(.+)$', content, re.MULTILINE)
title = title_match.group(1) if title_match else file_path.stem
# Get date from frontmatter or filename
date_str = fm_metadata.get('date')
if date_str and hasattr(date_str, 'strftime'):
date_str = date_str.strftime('%Y-%m-%d')
elif not date_str:
date_match = re.search(r'(\d{4}-\d{2}-\d{2})', file_path.name)
date_str = date_match.group(1) if date_match else "unknown"
# Merge frontmatter with file metadata
extra_metadata = {k: str(v) for k, v in fm_metadata.items()
if k not in ['title', 'date']}
else:
# Handle plain text files
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
title = file_path.stem
date_match = re.search(r'(\d{4}-\d{2}-\d{2})', file_path.name)
date_str = date_match.group(1) if date_match else "unknown"
extra_metadata = {}
# Create metadata
metadata = {
"title": title,
"filename": file_path.name,
"filepath": str(file_path),
"date": date_str,
"year": file_path.parent.parent.name if len(file_path.parts) > 2 else "unknown",
"month": file_path.parent.name if len(file_path.parts) > 1 else "unknown",
"file_size": len(content),
"created_at": datetime.now().isoformat(),
"file_extension": file_path.suffix
}
# Add any extra metadata from frontmatter
metadata.update(extra_metadata)
return {
"content": content,
"metadata": metadata
}
except Exception as e:
print(f"Error reading {file_path}: {e}")
return None
def chunk_content_smart(self, content: str, chunk_size: int = 1000, overlap: int = 200) -> List[str]:
"""Split content into overlapping chunks using sentence boundaries."""
if len(content) <= chunk_size:
return [content]
try:
# Use NLTK for better sentence splitting
sentences = sent_tokenize(content)
except Exception:
# Fallback to simple splitting if NLTK fails
sentences = re.split(r'[.!?]+\s+', content)
chunks = []
current_chunk = ""
for sentence in sentences:
# If adding this sentence would exceed chunk size
if len(current_chunk) + len(sentence) > chunk_size:
if current_chunk.strip():
chunks.append(current_chunk.strip())
# Start new chunk with overlap from previous chunk
words = current_chunk.split()
overlap_words = words[-overlap//10:] if len(words) > overlap//10 else words
current_chunk = " ".join(overlap_words) + " " + sentence
else:
# Single sentence is too long, split it
if len(sentence) > chunk_size:
words = sentence.split()
for i in range(0, len(words), chunk_size//10):
chunk_words = words[i:i + chunk_size//10]
chunks.append(" ".join(chunk_words))
else:
current_chunk = sentence
else:
current_chunk += " " + sentence if current_chunk else sentence
# Add the last chunk
if current_chunk.strip():
chunks.append(current_chunk.strip())
return chunks
def generate_chunk_id(self, filepath: str, chunk_content: str, chunk_index: int = 0) -> str:
"""Generate a unique ID for a document chunk based on content hash."""
content_hash = hashlib.md5(chunk_content.encode()).hexdigest()[:8]
path_hash = hashlib.md5(filepath.encode()).hexdigest()[:8]
return f"{path_hash}_{chunk_index}_{content_hash}"
def process_files_batch(self, batch_size: int = 50):
"""Process all supported files in the folder structure with batch insertion."""
if not self.output_folder.exists():
print(f"Output folder {self.output_folder} does not exist!")
return
processed_count = 0
skipped_count = 0
# Find all supported files
all_files = []
for ext in self.extensions:
all_files.extend(list(self.output_folder.glob(f"**/*{ext}")))
print(f"Found {len(all_files)} files to process...")
# Process files in batches
batch_documents = []
batch_metadatas = []
batch_ids = []
for file_path in all_files:
try:
# Extract content and metadata
doc_data = self.extract_content(file_path)
if not doc_data:
skipped_count += 1
continue
content = doc_data["content"]
metadata = doc_data["metadata"]
# Check if document already exists (by filepath)
existing = self.collection.get(
where={"filepath": str(file_path)}
)
if existing['ids']:
print(f"Skipping {file_path.name} (already processed)")
skipped_count += 1
continue
# Chunk the content
chunks = self.chunk_content_smart(content)
# Prepare batch data for each chunk
for i, chunk in enumerate(chunks):
doc_id = self.generate_chunk_id(str(file_path), chunk, i)
# Add chunk-specific metadata
chunk_metadata = metadata.copy()
chunk_metadata.update({
"chunk_index": i,
"total_chunks": len(chunks),
"chunk_size": len(chunk),
"chunk_hash": hashlib.md5(chunk.encode()).hexdigest()[:8],
"embedding_added_at": datetime.now().isoformat()
})
batch_documents.append(chunk)
batch_metadatas.append(chunk_metadata)
batch_ids.append(doc_id)
processed_count += 1
print(f"Prepared: {file_path.name} ({len(chunks)} chunks)")
# Insert batch when it reaches batch_size
if len(batch_documents) >= batch_size:
self.collection.add(
documents=batch_documents,
metadatas=batch_metadatas,
ids=batch_ids
)
print(f"Inserted batch of {len(batch_documents)} documents")
batch_documents.clear()
batch_metadatas.clear()
batch_ids.clear()
except Exception as e:
print(f"Error processing {file_path}: {e}")
skipped_count += 1
# Insert remaining documents
if batch_documents:
self.collection.add(
documents=batch_documents,
metadatas=batch_metadatas,
ids=batch_ids
)
print(f"Inserted final batch of {len(batch_documents)} documents")
print(f"\nProcessing complete!")
print(f"Processed: {processed_count} files")
print(f"Skipped: {skipped_count} files")
print(f"Total documents in collection: {self.collection.count()}")
def search_documents(self, query: str, n_results: int = 5, filter_metadata: Dict = None) -> List[Dict]:
"""Search for relevant documents using semantic similarity with optional filtering."""
try:
query_params = {
"query_texts": [query],
"n_results": n_results,
"include": ["documents", "metadatas", "distances"]
}
if filter_metadata:
query_params["where"] = filter_metadata
results = self.collection.query(**query_params)
search_results = []
for i in range(len(results['ids'][0])):
search_results.append({
'id': results['ids'][0][i],
'content': results['documents'][0][i],
'metadata': results['metadatas'][0][i],
'distance': results['distances'][0][i],
'relevance_score': 1 - results['distances'][0][i] # Convert distance to relevance
})
return search_results
except Exception as e:
print(f"Search error: {e}")
return []
def call_ollama(self, prompt: str, temperature: float = 0.7) -> str:
"""Call Ollama API with the given prompt."""
try:
response = requests.post(
f"{self.ollama_base_url}/api/generate",
json={
"model": self.model_name,
"prompt": prompt,
"stream": False,
"options": {
"temperature": temperature
}
},
timeout=120
)
if response.status_code == 200:
return response.json()["response"]
else:
return f"Error calling Ollama: {response.status_code} - {response.text}"
except requests.exceptions.RequestException as e:
return f"Error connecting to Ollama: {e}"
def answer_question(self, question: str, n_results: int = 5,
date_filter: str = None, min_relevance: float = 0.3) -> str:
"""Answer a question using retrieved context and Ollama with optional date filtering."""
print(f"Searching for: '{question}'")
# Prepare metadata filter
filter_metadata = {}
if date_filter:
# Simple date filtering - you could expand this
filter_metadata["date"] = date_filter
# Search for relevant documents
search_results = self.search_documents(question, n_results, filter_metadata)
if not search_results:
return "No relevant documents found for your question."
# Filter by relevance score
relevant_results = [r for r in search_results if r['relevance_score'] >= min_relevance]
if not relevant_results:
return f"No documents found with relevance score >= {min_relevance}. Try a different question or lower the threshold."
# Prepare context from search results
context_parts = []
for i, result in enumerate(relevant_results, 1):
metadata = result['metadata']
relevance = result['relevance_score']
context_parts.append(
f"Document {i} (from {metadata['filename']}, {metadata['date']}, relevance: {relevance:.2f}):\n"
f"{result['content']}\n"
)
context = "\n".join(context_parts)
# Create enhanced prompt for Ollama
prompt = f"""Based on the following documents, please answer the question accurately and comprehensively.
Use only the information provided in the documents. If the documents don't contain enough information to fully answer the question, please state that clearly.
Question: {question}
Relevant Documents:
{context}
Instructions:
- Provide a clear, well-structured answer
Answer: """
print("Generating answer with Ollama...")
answer = self.call_ollama(prompt)
# Add enhanced source information
sources = []
for result in relevant_results:
metadata = result['metadata']
relevance = result['relevance_score']
sources.append(f"- {metadata['filename']} ({metadata['date']}) - relevance: {relevance:.2f}")
return f"{answer}\n\nSources ({len(sources)} documents):\n" + "\n".join(sources)
def get_collection_stats(self) -> Dict:
"""Get statistics about the document collection."""
try:
count = self.collection.count()
if count == 0:
return {"total_documents": 0}
# Get sample of metadata to analyze
sample = self.collection.get(limit=min(100, count), include=["metadatas"])
dates = []
file_types = {}
years = {}
for metadata in sample['metadatas']:
# Count file types
ext = metadata.get('file_extension', 'unknown')
file_types[ext] = file_types.get(ext, 0) + 1
# Count years
year = metadata.get('year', 'unknown')
years[year] = years.get(year, 0) + 1
# Collect dates
date = metadata.get('date')
if date and date != 'unknown':
dates.append(date)
return {
"total_documents": count,
"file_types": file_types,
"years": years,
"date_range": {
"earliest": min(dates) if dates else "unknown",
"latest": max(dates) if dates else "unknown"
}
}
except Exception as e:
return {"error": str(e)}
def interactive_search(self):
"""Start an interactive search session with enhanced features."""
print("\n=== Interactive Markdown Search System ===")
print("Commands:")
print(" - Type your question to search")
print(" - 'stats' to see collection statistics")
print(" - 'help' to see this help")
print(" - 'quit' to exit")
print("-" * 50)
# Show initial stats
stats = self.get_collection_stats()
if 'error' not in stats:
print(f"Collection loaded: {stats['total_documents']} documents")
if stats['total_documents'] > 0:
print(f"Date range: {stats['date_range']['earliest']} to {stats['date_range']['latest']}")
print(f"File types: {', '.join(f'{k}({v})' for k, v in stats['file_types'].items())}")
print("-" * 50)
while True:
try:
user_input = input("\n> ").strip()
if user_input.lower() in ['quit', 'exit', 'q']:
print("Goodbye!")
break
if user_input.lower() == 'help':
print("\nCommands:")
print(" - Ask any question about your documents")
print(" - 'stats' - Show collection statistics")
print(" - 'quit' - Exit the program")
continue
if user_input.lower() == 'stats':
stats = self.get_collection_stats()
print(f"\nCollection Statistics:")
for key, value in stats.items():
print(f" {key}: {value}")
continue
if not user_input:
continue
print("\nSearching and generating answer...")
answer = self.answer_question(user_input)
print(f"\nAnswer:\n{answer}")
print("-" * 50)
except KeyboardInterrupt:
print("\nGoodbye!")
break
except Exception as e:
print(f"Error: {e}")
def agent_strategy_factory(search_system: MarkdownSearchSystem, context_hint: str = "") -> Callable[[str], str]:
def strategy(query: str) -> str:
refined_query = f"{context_hint} {query}".strip()
return search_system.answer_question(refined_query, n_results=7)
return strategy
def main():
"""Main function to run the search system."""
import argparse
parser = argparse.ArgumentParser(description="Enhanced Markdown Search System with ChromaDB and Ollama")
parser.add_argument("--output-folder", default="output", help="Path to output folder containing files")
parser.add_argument("--db-path", default="./chroma_db", help="Path for ChromaDB storage")
parser.add_argument("--collection", default="markdown_docs", help="ChromaDB collection name")
parser.add_argument("--model", default="mistral", help="Ollama model name")
parser.add_argument("--embedding-model", default="nomic-embed-text", help="Ollama embedding model")
parser.add_argument("--extensions", nargs='+', default=[".md", ".txt"], help="File extensions to process")
parser.add_argument("--batch-size", type=int, default=50, help="Batch size for processing")
parser.add_argument("--process", action="store_true", help="Process files into ChromaDB")
parser.add_argument("--search", action="store_true", help="Start interactive search")
parser.add_argument("--question", help="Ask a single question")
parser.add_argument("--stats", action="store_true", help="Show collection statistics")
parser.add_argument("--reset-collection", action="store_true", help="Delete and recreate the collection")
args = parser.parse_args()
# Handle collection reset
if args.reset_collection:
print(f"Resetting collection '{args.collection}'...")
try:
client = chromadb.PersistentClient(path=args.db_path)
try:
client.delete_collection(name=args.collection)
print(f"Deleted existing collection: {args.collection}")
except Exception as e:
print(f"Collection didn't exist or couldn't be deleted: {e}")
except Exception as e:
print(f"Error connecting to ChromaDB: {e}")
return
# Initialize the search system
search_system = MarkdownSearchSystem(
output_folder=args.output_folder,
db_path=args.db_path,
collection_name=args.collection,
model_name=args.model,
embedding_model=args.embedding_model,
extensions=args.extensions
)
if args.process:
print("Processing files...")
search_system.process_files_batch(batch_size=args.batch_size)
elif args.stats:
stats = search_system.get_collection_stats()
print("\nCollection Statistics:")
for key, value in stats.items():
print(f" {key}: {value}")
elif args.question:
answer = search_system.answer_question(args.question)
print(f"\nQuestion: {args.question}")
print(f"Answer:\n{answer}")
# Example usage inside main()
# Replace the `elif args.search:` block with the following
elif args.search:
print("Agentic search enabled")
# Create agents with slightly different strategies or focuses
agent1 = ResearchAgent(
name="Contextualizer",
goal=args.question or "Summarize the main topics in my markdown corpus.",
strategy=agent_strategy_factory(search_system, context_hint="Extract context and definitions")
)
agent2 = ResearchAgent(
name="Synthesizer",
goal=agent1.goal,
strategy=agent_strategy_factory(search_system, context_hint="Synthesize and summarize")
)
orchestrator = ResearchOrchestrator(
agents=[agent1, agent2],
goal=agent1.goal
)
orchestrator.run()
else:
print("Enhanced Markdown Search System")
print("\nUsage examples:")
print(" python analyze.py --process # Process files into ChromaDB")
print(" python analyze.py --search # Interactive search")
print(" python analyze.py --stats # Show collection stats")
print(" python analyze.py --question 'What happened in January 2024?'")
print(" python analyze.py --process --search # Process then search")
print(" python analyze.py --extensions .md .txt .rst # Process multiple file types")
print(" python analyze.py --reset-collection # Delete existing collection")
print(" python analyze.py --reset-collection --process # Reset and reprocess")
if __name__ == "__main__":
main()
Thanks! From your message, I see you're asking: **how can we make the agents dynamically propose new queries based on partial answers?**
Here's how we can add that:
---
### 🧠 Goal:
Allow each agent to read partial responses and generate **refined or follow-up queries** based on what’s missing or ambiguous.
---
### ✅ Step-by-Step Plan
1. **Add a `Memory` or `ContextHistory` object**
Store all previous messages, answers, and queries.
2. **Let each `ResearchAgent` have a `think()` method**
This method takes the context, thinks about what’s missing, and outputs a new query.
3. **Use `ollama` or a local LLM to generate refinements**
For example:
```python
def think(self, previous_results: List[str]) -> str:
prompt = f"Here is the current answer:\n{previous_results[-1]}\n\nWhat question should be asked next to clarify or expand?"
response = self.llm(prompt)
return response.strip()
```
4. **Loop until convergence or limit**
Track iterations and stop when:
- Confidence is high (e.g. LLM says “this is sufficient”), or
- No new meaningful queries are generated.
---
### 🧪 Example Update to `ResearchAgent`
```python
class ResearchAgent:
def __init__(self, name, strategy_fn, model=None):
self.name = name
self.strategy = strategy_fn
self.model = model # Optional: e.g. Ollama or local LLM
self.history = []
def think(self, last_answer: str, question: str) -> str:
prompt = (
f"Given the original question:\n'{question}'\n"
f"And this partial answer:\n'{last_answer}'\n"
"What follow-up or clarifying question should be asked next?"
)
response = self.model(prompt)
return response.strip()
def act(self, system: MarkdownSearchSystem, query: str):
answer = self.strategy(system, query)
self.history.append((query, answer))
return answer
```
---
### 🔁 Update to `ResearchOrchestrator.run()`
Add logic like:
```python
for agent in self.agents:
new_query = agent.think(context, question)
response = agent.act(system, new_query)
context += f"\n\nAgent {agent.name} follow-up: {new_query}\n{response}\n"
```
---
### ⚡️ Bonus Ideas
- You can use `nomic-embed-text` to embed each partial answer and compute cosine similarity between iterations to detect stagnation.
- Log outputs to a Markdown `session_log.md` for review.
- Use a final summarizing agent (`ReporterAgent`) to generate a human-readable report.
---
Want a full version of this implemented in the current script? I can append it to the `analyze.py` you're building.
gpt-4o