| """RAG Agent using LlamaIndex for codebase analysis.""" |
|
|
| import os |
| from typing import Dict, List, Optional |
| from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings |
| from llama_index.llms.openai import OpenAI |
| from llama_index.core.node_parser import SimpleNodeParser |
|
|
| class RAGAgent: |
| """Agent for RAG-based codebase analysis using LlamaIndex.""" |
|
|
| def __init__(self): |
| self.index = None |
| self.query_engine = None |
| self._setup_llm() |
|
|
| def _setup_llm(self): |
| """Configure LlamaIndex settings.""" |
| api_key = os.getenv("OPENAI_API_KEY") |
| if api_key: |
| Settings.llm = OpenAI(model="gpt-4o-mini", api_key=api_key) |
|
|
| def index_codebase(self, directory_path: str) -> str: |
| """Index the codebase directory.""" |
| if not os.getenv("OPENAI_API_KEY"): |
| return "❌ OpenAI API Key required for LlamaIndex RAG." |
|
|
| try: |
| |
| reader = SimpleDirectoryReader( |
| input_dir=directory_path, |
| recursive=True, |
| exclude_hidden=True, |
| required_exts=[".py", ".js", ".ts", ".md", ".json", ".yml", ".yaml"] |
| ) |
| documents = reader.load_data() |
|
|
| |
| self.index = VectorStoreIndex.from_documents(documents) |
| self.query_engine = self.index.as_query_engine() |
| |
| return f"✅ Indexed {len(documents)} files successfully." |
| except Exception as e: |
| return f"❌ Indexing failed: {str(e)}" |
|
|
| def query_codebase(self, query: str) -> str: |
| """Query the indexed codebase.""" |
| if not self.query_engine: |
| return "⚠️ Codebase not indexed. Please analyze a folder first." |
| |
| try: |
| response = self.query_engine.query(query) |
| return str(response) |
| except Exception as e: |
| return f"❌ Query failed: {str(e)}" |
|
|
