Spaces:
Sleeping
Sleeping
File size: 3,827 Bytes
eb5a3ff | 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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | 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
})
|