agent.py

from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_community.llms import OpenAI as LangOpenAI
from tools.search import web_search

def load_llm():
    return LangOpenAI(temperature=0, model_name="gpt-4")

def build_prompt():
    template = """You are a helpful research assistant. Answer the question using the context provided.

Context:
{context}

Question: {question}
Answer:"""
    return PromptTemplate(template=template, input_variables=["question", "context"])

def answer_question(question, file_context=None, do_search=False):
    context_parts = []
    if file_context:
        context_parts.append(file_context)
    if do_search:
        context_parts.append(web_search(question))

    context = "\n\n".join(context_parts) if context_parts else "N/A"

    llm = load_llm()
    chain = LLMChain(llm=llm, prompt=build_prompt())
    return chain.run({"question": question, "context": context}).strip()