Agentic RAG: Enhancing retrieval-augmented generation with AI agents

import os
import glob
import re
import requests
import time
from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma

# GitHub repository details
REPOS = {
    "weave": "wandb/weave",
    "wandb": "wandb/docs",  # For documentation
}

# GitHub repositories for issues
ISSUE_REPOS = {
    "weave": "wandb/weave",
    "wandb": "wandb/wandb",  # Correct repo for WandB issues
}

# GitHub issues URL
ISSUES_ENDPOINT = "/issues?state=all&per_page=100"

# Initialize OpenAI Embeddings
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# Function to download repo docs using manual git commands
def download_docs(repo_name, repo_url, path):
    if os.path.exists(path):
        print(f"{repo_name} documentation already exists, skipping download.")
        return
    
    print(f"Downloading {repo_name} documentation...")
    
    # Simple git clone without sparse-checkout (more reliable)
    clone_cmd = f"git clone --depth 1 https://github.com/{repo_url}.git {path}"
    result = os.system(clone_cmd)
    
    if result != 0 or not os.path.exists(path):
        print(f"❌ Failed to clone {repo_name}. Check if the repository URL is correct.")
        return
    
    print(f"✅ {repo_name} documentation downloaded successfully.")

# Function to clean markdown content
def clean_markdown(text):
    if not text:
        return ""
    text = re.sub(r'!\[.*?\]\(.*?\)', '', text)  # Remove images
    text = re.sub(r'\[([^\]]+)\]\((.*?)\)', r'\1', text)  # Keep link text but remove URLs
    return text.strip()

# Function to fetch GitHub issues without comments
def fetch_issues(repo_name, repo_url, max_pages=5):
    all_issues = []
    page = 1
    
    # Paginate through the results
    while page <= max_pages:
        url = f"https://api.github.com/repos/{repo_url}{ISSUES_ENDPOINT}&page={page}"
        headers = {"Accept": "application/vnd.github.v3+json"}
        
        try:
            print(f"Fetching page {page} of issues from {repo_name}...")
            response = requests.get(url, headers=headers)
            
            if response.status_code == 200:
                issues = response.json()
                
                # If no more issues are returned, break the loop
                if not issues:
                    break
                
                # Process each issue without fetching comments
                for issue in issues:
                    # Add basic issue data to the list
                    all_issues.append({
                        "id": issue["id"],
                        "number": issue["number"],
                        "title": issue["title"],
                        "body": issue.get("body", ""),
                        "state": issue["state"],
                        "created_at": issue["created_at"],
                        "comments_count": issue.get("comments", 0),
                        "github_url": issue["html_url"],  # URL to the GitHub issue page
                        "comments_url": issue["comments_url"],  # Direct API URL for comments
                        "repo": repo_url  # Store the repo info for later reference
                    })
                
                print(f"Fetched {len(issues)} issues from {repo_name} (page {page})")
                
                page += 1
                time.sleep(1)  # Small delay between pages
                
            elif response.status_code == 403 and "rate limit exceeded" in response.text.lower():
                print("⚠️ GitHub API rate limit exceeded. Waiting for 60 seconds...")
                time.sleep(60)  # Wait for rate limit to reset
                continue  # Retry the same page
            else:
                print(f"Error fetching issues from {repo_name}: {response.status_code}")
                break
                
        except Exception as e:
            print(f"Exception fetching issues for {repo_name}: {e}")
            break
    
    print(f"Total issues fetched from {repo_name}: {len(all_issues)}")
    return all_issues

# Function to create documents from issues (without comments)
def create_issue_documents(issues):
    docs = []
    for issue in issues:
        # Create a document for the main issue
        issue_text = f"Title: {issue['title']}\nIssue #{issue['number']}\nState: {issue['state']}\nComments: {issue['comments_count']}\n"
        
        # Add a note about comments and where to find them
        if issue['comments_count'] > 0:
            issue_text += f"\nThis issue has {issue['comments_count']} comments. View at: {issue['github_url']}\n"
        
        issue_text += f"\nBody:\n{issue['body']}"
        
        # Create the document object with rich metadata
        docs.append(Document(
            page_content=clean_markdown(issue_text),
            metadata={
                "id": issue["id"],
                "number": issue["number"],
                "created_at": issue["created_at"],
                "state": issue["state"],
                "comments_count": issue["comments_count"],
                "github_url": issue["github_url"],
                "comments_url": issue["comments_url"],
                "repo": issue["repo"],
                "type": "issue"
            }
        ))
    
    return docs

# Function to process Markdown docs
def create_markdown_documents(repo_name, doc_path):
    if not os.path.exists(doc_path):
        print(f"⚠️ {doc_path} does not exist, skipping doc processing.")
        return []
    
    # Find markdown files recursively
    md_files = glob.glob(f'{doc_path}/**/*.md', recursive=True)
    
    if not md_files:
        print(f"No markdown files found in {doc_path}.")
        return []
    
    docs = []
    for md_file in md_files:
        try:
            with open(md_file, 'r', encoding='utf-8') as f:
                content = clean_markdown(f.read())
                # Create a proper Langchain Document object
                docs.append(Document(
                    page_content=content,
                    metadata={
                        "file": md_file,
                        "type": "documentation"
                    }
                ))
        except Exception as e:
            print(f"Error processing {md_file}: {e}")
    
    return docs

# Function to query the vector store
def search_docs(query, db_directory, k=5):
    """Search a Chroma vector store for documents similar to the query."""
    try:
        # Load the Chroma database
        db = Chroma(persist_directory=db_directory, embedding_function=embeddings)
        
        # Search for similar documents
        results = db.similarity_search_with_score(query, k=k)
        
        if not results:
            print(f"No results found for '{query}'")
            return []
        
        print(f"Found {len(results)} results for '{query}':")
        
        for i, (doc, score) in enumerate(results):
            print(f"\nResult {i+1}:")
            print(f"Score: {score}")
            
            # Print metadata
            print("Metadata:")
            for key, value in doc.metadata.items():
                if key in ["id", "number", "state", "comments_count", "type"]:
                    print(f"  {key}: {value}")
            
            # For issues with comments, show the GitHub URL
            if doc.metadata.get("type") == "issue" and doc.metadata.get("comments_count", 0) > 0:
                print(f"  GitHub URL: {doc.metadata.get('github_url')}")
            
            # Print content (truncated if too long)
            content = doc.page_content
            if len(content) > 500:
                content = content[:500] + "... [truncated]"
            print(f"\nContent: {content}\n")
            print("-" * 40)
        
        return results
    
    except Exception as e:
        print(f"Error searching database: {e}")
        return []

# Main execution
if __name__ == "__main__":
    # Create directory for vector stores
    os.makedirs("db", exist_ok=True)
    
    for repo_name, repo_url in REPOS.items():
        print(f"\nProcessing {repo_name.upper()}...")

        # Fetch issues from the correct repository
        issue_repo_url = ISSUE_REPOS[repo_name]
        print(f"Fetching issues from {issue_repo_url}...")
        issues = fetch_issues(repo_name, issue_repo_url)

        # Download docs
        doc_path = f"{repo_name}_docs"
        download_docs(repo_name, repo_url, doc_path)

        # Process issues
        if issues:
            print(f"Creating vector store for {repo_name} issues...")
            issue_docs = create_issue_documents(issues)
            
            # Create Chroma vector store for issues
            issues_db = Chroma.from_documents(
                documents=issue_docs,
                embedding=embeddings,
                persist_directory=f"db/{repo_name}_issues"
            )
            
            print(f"✅ Saved {repo_name} issues vector store ({len(issue_docs)} documents)")
        else:
            print("No issues to process.")

        # Process docs
        print(f"Creating vector store for {repo_name} documentation...")
        md_docs = create_markdown_documents(repo_name, doc_path)
        
        if md_docs:
            # Create Chroma vector store for documentation
            docs_db = Chroma.from_documents(
                documents=md_docs,
                embedding=embeddings,
                persist_directory=f"db/{repo_name}_docs"
            )
            
            print(f"✅ Saved {repo_name} documentation vector store ({len(md_docs)} documents)")
        else:
            print("No documentation to process.")

    print("\n✅ Data ingestion completed successfully.")
    
    # Example search to verify everything works
    print("\nTesting search functionality...")
    search_docs("authentication issues", "db/weave_issues", k=2)

You are a query-routing assistant that determines whether a user's question is related to Weights & Biases (wandb) or Weave and directs them to the most relevant database. Users may ask about documentation, troubleshooting issues, or implementation details for either tool.

Your task is to analyze the query and route it to all relevant sources if uncertainty exists. If a query could apply to multiple areas, return multiple results rather than choosing just one.

The available knowledge bases are:
1. "Weave Issues" - Contains GitHub issues from the wandb/weave repository
2. "Weave Documentation" - Contains documentation for Weave
3. "WandB Issues" - Contains GitHub issues from the wandb/wandb repository
4. "WandB Documentation" - Contains documentation for Weights & Biases

Return a JSON object with the following format:
{
  "knowledge_bases": ["name_of_kb1", "name_of_kb2"]
}

## Library Descriptions for Context

### Weights & Biases (wandb) – Experiment Tracking & ML Workflow Management
Weights & Biases (wandb) is a tool for tracking and visualizing machine learning experiments. It helps researchers and engineers:
* Log and monitor metrics like loss, accuracy, and learning rate.
* Track hyperparameters and model configurations for experiment reproducibility.
* Store and visualize model artifacts and training progress.
* Integrate seamlessly with deep learning frameworks like PyTorch, TensorFlow, and Hugging Face Transformers.
* Collaborate via dashboards that display real-time and historical training runs.
It is widely used in research and production to optimize models, compare runs, and fine-tune hyperparameters efficiently.

### W&B Weave – Generative AI & LLM Debugging Toolkit
Weave is a toolkit designed for developing, debugging, and evaluating Generative AI applications, particularly those using Large Language Models (LLMs). It enables users to:
* Log inputs and outputs from LLMs for better debugging.
* Build evaluations to compare different model responses.
* Organize information across LLM workflows from experimentation to production.
* Track interactions within multi-step AI pipelines to analyze decision-making processes.
Weave is essential for understanding and improving the performance of AI agents, LLM-based chatbots, and retrieval-augmented generation (RAG) systems.

## Routing Logic
* If the user asks about experiment tracking, ML metrics, logging, hyperparameters, dashboards, or integrations, route to W&B Docs.
* If they report an error, unexpected behavior, or API issue related to W&B, route to W&B Issues.
* If the user asks about LLM evaluation, function tracking, Weave UI development, or Weave's role in ML workflows, route to Weave Docs.
* If they report errors related to Weave tracking or data storage, route to Weave Issues.
* If uncertain, include all possibly relevant sources instead of picking only one.

## Code Samples for Context
If the query is vague or unclear, match it against the functionality of each tool using these sample code snippets:

### W&B Experiment Tracking Example (wandb Docs/Issues)
```python
import wandb
wandb.init(project="my_project", config={"learning_rate": 0.001, "epochs": 10})
for epoch in range(wandb.config.epochs):
    train_loss = 0.01 * (wandb.config.epochs - epoch)
    train_accuracy = 0.1 * epoch
    wandb.log({'epoch': epoch, 'train_loss': train_loss, 'train_accuracy': train_accuracy})
```
If the query involves logging, training loops, metrics, dashboards, or hyperparameters, it should likely be routed to W&B Docs or W&B Issues if errors are involved.

### Weave Function Tracking Example (Weave Docs/Issues)
```python
import weave
weave.init(project_name="my_weave_project")
@weave.op()
def add_numbers(a: int, b: int) -> int:
    return a + b
result = add_numbers(5, 7)
print(f'Result: {result}')
```
If the query involves tracking function calls, logging LLM interactions, debugging AI workflows, or analyzing execution steps, route it to Weave Docs or Weave Issues if errors occur.


Focus on understanding:
- Is the query about Weave or WandB?
- Is the user looking for documentation or issues/bugs?
- If the query mentions "bug" or "issue" or describes a problem, prioritize the Issues knowledge base
- If the query is asking how to use a feature or understand concepts, prioritize Documentation
- If unclear whether the query is about Weave or WandB, include knowledge bases for both

Examples:
1. For "How do I track experiments in wandb?"
   - Knowledge bases: ["WandB Documentation"]

2. For "Weave authentication failing"
   - Knowledge bases: ["Weave Issues"]

3. For "How to visualize neural networks in Weights & Biases"
   - Knowledge bases: ["WandB Documentation"]

4. For "LLM output tracking shows errors in production"
   - Knowledge bases: ["Weave Issues", "Weave Documentation"]

Always include at least one knowledge base in your response.

Return a JSON object with the following format:
{
  "knowledge_bases": ["name_of_kb1", "name_of_kb2"]
}

import sys
import os
import json
from litellm import completion
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
import weave; weave.init("agentic_rag")
import requests
import os
import json
import time

# Create cache directory if it doesn't exist
os.makedirs("comment_cache", exist_ok=True)

# Initialize OpenAI Embeddings
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# Configure which databases to search (can be modified to include/exclude databases)
DATABASES_TO_SEARCH = [
    {"name": "Weave Issues", "path": "db/weave_issues"},
    {"name": "Weave Documentation", "path": "db/weave_docs"},
    {"name": "WandB Issues", "path": "db/wandb_issues"},
    {"name": "WandB Documentation", "path": "db/wandb_docs"}
]

def load_router_prompt():
    """Load the router prompt from router_prompt.txt"""
    try:
        with open("router_prompt.txt", "r") as f:
            return f.read()
    except FileNotFoundError:
        print("❌ Error: router_prompt.txt not found.")
        print("Please create this file with your router prompt before running the script.")
        sys.exit(1)

def search_vector_store(query, db_directory, k=5):
    """Search a vector store and return results."""
    try:
        # Load the Chroma database
        db = Chroma(persist_directory=db_directory, embedding_function=embeddings)
        
        # Search for similar documents
        results = db.similarity_search_with_score(query, k=k)
        
        return results
    
    except Exception as e:
        print(f"❌ Error searching database at {db_directory}: {e}")
        return []


def extract_json_from_response(response_text):
    """Extract JSON from response text, handling code blocks."""
    # Try to find JSON in code blocks
    import re
    json_pattern = r'```(?:json)?\s*(\{.*?\})\s*```'
    match = re.search(json_pattern, response_text, re.DOTALL)
    
    if match:
        # Extract JSON from code block
        json_str = match.group(1)
    else:
        # Assume the entire response is JSON
        json_str = response_text
    
    # Remove any non-JSON content
    json_str = json_str.strip()
    
    return json.loads(json_str)


def route_query(user_query):
    """Use LiteLLM to route the query to the appropriate databases."""
    router_prompt = load_router_prompt()
    
    try:
        # Call LiteLLM with the router prompt
        response = completion(
            model="openai/gpt-4o",
            messages=[
                {"role": "system", "content": router_prompt},
                {"role": "user", "content": user_query}
            ]
        )
        
        # Extract the content from the response
        router_response = response.choices[0].message.content
        
        # Parse the JSON response
        try:
            # parsed_response = json.loads(router_response)
            parsed_response = extract_json_from_response(router_response)
            return parsed_response
        except json.JSONDecodeError:
            print("❌ Error: Router didn't return valid JSON. Using all knowledge bases.")
            print(f"Router response: {router_response}")
            return {
                "knowledge_bases": [db["name"] for db in DATABASES_TO_SEARCH],
                "search_query": user_query,
                "reasoning": "Fallback due to JSON parsing error"
            }
    
    except Exception as e:
        print(f"❌ Error calling LiteLLM: {e}")
        return {
            "knowledge_bases": [db["name"] for db in DATABASES_TO_SEARCH],
            "search_query": user_query,
            "reasoning": f"Fallback due to error: {str(e)}"
        }

def get_kb_by_name(name):
    """Get the knowledge base by name."""
    for db in DATABASES_TO_SEARCH:
        if db["name"] == name:
            return db
    return None

def fetch_comments_for_issue(repo, issue_number):

    # Check cache first
    cache_file = f"comment_cache/{repo.replace('/', '_')}_issue_{issue_number}_comments.json"
    if os.path.exists(cache_file):
        try:
            with open(cache_file, 'r') as f:
                comments = json.load(f)
                print(f"Loaded {len(comments)} comments for issue #{issue_number} from cache")
                return comments
        except Exception as e:
            print(f"Error reading cache: {e}")
            # Continue to fetch if cache read fails
    
    # Fetch from GitHub API
    url = f"https://api.github.com/repos/{repo}/issues/{issue_number}/comments"
    headers = {"Accept": "application/vnd.github.v3+json"}
    
    # Add GitHub token if available
    if os.environ.get("GITHUB_TOKEN"):
        headers["Authorization"] = f"token {os.environ.get('GITHUB_TOKEN')}"
    
    try:
        print(f"Fetching comments for issue #{issue_number} from {repo}...")
        response = requests.get(url, headers=headers)
        
        if response.status_code == 200:
            comments = response.json()
            
            # Cache the comments
            with open(cache_file, 'w') as f:
                json.dump(comments, f, indent=2)
            
            print(f"✅ Fetched and cached {len(comments)} comments for issue #{issue_number}")
            return comments
            
        elif response.status_code == 403 and "rate limit exceeded" in response.text.lower():
            print("⚠️ GitHub API rate limit exceeded. Waiting 60 seconds...")
            time.sleep(60)
            return fetch_comments_for_issue(repo, issue_number)  # Retry after waiting
            
        else:
            print(f"Error fetching comments: {response.status_code}")
            return []
            
    except Exception as e:
        print(f"Exception fetching comments: {e}")
        return []

def format_comments_for_context(comments):

    if not comments:
        return ""
    
    comments_text = "\n\nCOMMENTS:\n"
    for i, comment in enumerate(comments):
        user = comment.get('user', {}).get('login', 'anonymous')
        created_at = comment.get('created_at', 'unknown date')
        body = comment.get('body', '').strip()
        
        comments_text += f"\n[Comment #{i+1} by {user} on {created_at}]\n{body}\n"
    
    return comments_text

def generate_response(user_query, relevant_docs, query_analysis):
    """Generate a response using LiteLLM based on the retrieved documents."""
    # Prepare context from documents
    context = ""
    github_links = []
    
    for i, (doc, score) in enumerate(relevant_docs):
        # Add the document to context
        context += f"--- Document {i+1} (Relevance: {score:.4f}) ---\n"
        
        # Add source information
        if doc.metadata.get('type') == 'issue':
            issue_num = doc.metadata.get('number', 'Unknown')
            repo = doc.metadata.get('repo', 'Unknown')
            context += f"Source: GitHub Issue #{issue_num} in {repo}\n"
            
            # If issue has comments, fetch and add them to context
            if doc.metadata.get('comments_count', 0) > 0:
                github_url = doc.metadata.get('github_url')
                if github_url:
                    github_links.append(github_url)
                
                # Fetch and append comments
                print(f"Fetching comments for issue #{issue_num}...")
                comments = fetch_comments_for_issue(repo, issue_num)
                if comments:
                    comments_text = format_comments_for_context(comments)
                    context += doc.page_content + comments_text + "\n\n"
                    continue  # Skip the standard content addition below
        else:
            context += f"Source: {doc.metadata.get('file', 'Unknown')}\n"
        
        # Add content (only reached if the issue doesn't have comments or isn't an issue)
        context += doc.page_content
        context += "\n\n"
    
    # Create the prompt for the response
    response_prompt = f"""You are an AI assistant for Weights & Biases (wandb) and Weave.
You've been given several documents retrieved from a search based on the user's query.
Use the information in these documents to answer the user's question.

If the documents don't contain the necessary information to answer the question, admit that you don't know
rather than making up an answer. If appropriate, suggest what the user might search for instead.

User Query: {user_query}

Retrieved Documents:
{context}

Based on these documents, provide a helpful response to the user's query.
Pay special attention to issue comments as they often contain solutions to the problems described in the issues.
"""

    try:
        # Call LiteLLM with the response prompt
        response = completion(
            model="openai/gpt-4o",
            messages=[
                {"role": "system", "content": response_prompt},
                {"role": "user", "content": user_query}
            ]
        )
        
        # Extract the content from the response
        return response.choices[0].message.content
    
    except Exception as e:
        print(f"❌ Error generating response: {e}")
        return f"I encountered an error while generating a response: {str(e)}"


def agentic_rag(user_query):
    """Main function for the agentic RAG system."""
    print(f"\n🔍 Processing query: '{user_query}'")
    print("-" * 80)
    
    # Step 1: Route the query
    query_analysis = route_query(user_query)
    print(f"✅ Query Analysis:")
    print(f"- Selected Knowledge Bases: {', '.join(query_analysis['knowledge_bases'])}")
    # print(f"- Refined Search Query: '{query_analysis['search_query']}'")
    # print(f"- Reasoning: {query_analysis['reasoning']}")
    
    # Step 2: Search the selected knowledge bases
    all_results = []
    
    for kb_name in query_analysis["knowledge_bases"]:
        kb = get_kb_by_name(kb_name)
        if kb and os.path.exists(kb["path"]):
            # print(f"\n🔍 Searching '{kb['name']}' for: '{query_analysis['search_query']}'")
            results = search_vector_store(user_query, kb["path"], k=3)
            
            if results:
                print(f"Found {len(results)} results in {kb['name']}")
                
                # Print out GitHub URLs for issues with comments
                for doc, _ in results:
                    if doc.metadata.get('type') == 'issue' and doc.metadata.get('comments_count', 0) > 0:
                        github_url = doc.metadata.get('github_url')
                        if github_url:
                            print(f"  - Issue with comments: {github_url}")
                
                all_results.extend(results)
            else:
                print(f"No results found in {kb['name']}")
        else:
            print(f"⚠️ Knowledge base '{kb_name}' not found or not available.")
    
    # Step 3: Sort results by score (ascending since lower score = more relevant in Chroma)
    all_results.sort(key=lambda x: x[1])
    
    # Take top results (up to 5)
    top_results = all_results[:5]
    
    if not top_results:
        print("❌ No relevant documents found across any knowledge base.")
        return "I couldn't find any relevant information to answer your question. Could you please rephrase or provide more details?"
    
    # Step 4: Generate a response
    print(f"\n💬 Generating response based on {len(top_results)} documents...")
    response = generate_response(user_query, top_results, query_analysis)
    
    return response

def interactive_mode():
    """Run the RAG system in interactive mode."""
    print("\n" + "=" * 80)
    print("AGENTIC RAG SYSTEM")
    print("=" * 80)
    print("Type 'exit' to quit.")
    print("To ask a question, type your question (can be multiple lines)")
    print("When finished, press Enter on an empty line to submit.")
    
    while True:
        print("\nYour question (multi-line, empty line to submit):")
        lines = []
        
        while True:
            line = input().strip()
            
            # Check for exit command on single line
            if not lines and line.lower() in ['exit', 'quit']:
                return
            
            # Empty line finishes input if we already have some content
            if not line and lines:
                break
                
            # Otherwise add the line to our input
            if line:
                lines.append(line)
        
        # Combine all lines into a single query
        user_query = "\n".join(lines)
        
        if not user_query:
            print("Please enter a question.")
            continue
        
        print("\n" + "-" * 80)
        print("Processing your query...")
        
        response = agentic_rag(user_query)
        
        print("\n" + "=" * 80)
        print("ANSWER:")
        print(response)
        print("=" * 80)

if __name__ == "__main__":
    # Check for API key
    if not os.environ.get("OPENAI_API_KEY"):
        api_key = input("Please enter your OpenAI API key: ").strip()
        if api_key:
            os.environ["OPENAI_API_KEY"] = api_key
        else:
            print("❌ No API key provided. Exiting.")
            sys.exit(1)
    
    interactive_mode()