agent.py

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.documents import Document
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter

load_dotenv()

def get_llm():
    provider = os.getenv("LLM_PROVIDER", "google").lower()
    
    if provider == "openai":
        return ChatOpenAI(
            model=os.getenv("OPENAI_MODEL", "gpt-4.1-mini"), 
            temperature=0.7
        )
    elif provider == "google":
        return ChatGoogleGenerativeAI(
            model=os.getenv("GOOGLE_MODEL", "gemini-3.1-flash-lite-preview"),
            temperature=0.7
        )
    else:
        raise ValueError(f"Unknown LLM Provider: {provider}")

def get_embeddings():
    provider = os.getenv("LLM_PROVIDER", "google").lower()
    if provider == "openai":
        return OpenAIEmbeddings(model="text-embedding-3-small")
    elif provider == "google":
        return GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001")
    else:
        raise ValueError(f"Unknown LLM Provider: {provider}")

def process_and_retrieve_context(description: str, field: str, files_data: list[dict]) -> str:
    """Takes a list of file dictionaries and retrieves relevant context using ChromaDB."""
    if not files_data:
        return "No extra files provided."
        
    docs = []
    for file in files_data:
        docs.append(Document(
            page_content=file["content"],
            metadata={"source": file["filename"]}
        ))
        
    # Split the documents
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000, 
        chunk_overlap=200
    )
    splits = text_splitter.split_documents(docs)
    
    # Store locally in chromadb directory and use it to retrieve
    vectorstore = Chroma.from_documents(
        documents=splits, 
        embedding=get_embeddings(),
        persist_directory="./chroma_db"
    )
    
    # Use description and field to retrieve relevant chunks
    query = f"Field: {field}. Task: {description}"
    retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
    
    retrieved_docs = retriever.invoke(query)
    
    context = ""
    for idx, doc in enumerate(retrieved_docs):
        context += f"\n--- Retrieved Chunk {idx+1} from {doc.metadata.get('source', 'Unknown')} ---\n{doc.page_content}\n"
        
    return context

def generate_task_prompt(description: str, field: str, files_data: list[dict]) -> str:
    llm = get_llm()
    
    # Get filtered context via RAG
    files_context = process_and_retrieve_context(description, field, files_data)
    
    system_prompt = (
        "You are an expert technical project manager and architect. "
        "Your goal is to take a task description provided by a project manager, context about the field (e.g., backend, frontend), "
        "and any uploaded file context, and produce a high-quality, developer-ready task prompt.\n\n"
        "Return ONLY the finalized prompt ready to be handed to a developer."
    )
    
    human_prompt = (
        "Field/Domain: {field}\n"
        "Task Description:\n{description}\n\n"
        "Relevant Code/Files Context:\n{files_context}\n\n"
        "Please generate a comprehensive developer prompt."
    )
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", human_prompt),
    ])
    
    chain = prompt | llm | StrOutputParser()
    
    return chain.invoke({
        "field": field,
        "description": description,
        "files_context": files_context
    })