File size: 1,966 Bytes
a8063f2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
"""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:
# Load documents
reader = SimpleDirectoryReader(
input_dir=directory_path,
recursive=True,
exclude_hidden=True,
required_exts=[".py", ".js", ".ts", ".md", ".json", ".yml", ".yaml"]
)
documents = reader.load_data()
# Create index
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)}"
|