agents/citation_formatter.py

"""
LangChain LLMChain for generating bibliography from research notes and draft.
"""
from langchain.callbacks.base import BaseCallbackHandler
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_openai import OpenAI

import os

# Prompt template for bibliography generation
BIBLIOGRAPHY_PROMPT = """
You are an expert academic writer. Given research notes and a draft, generate a bibliography in APA format with a maximum of 10 references.

Research notes:
{research_notes}

Draft:
{draft}

Instructions:
- Analyze the research notes and draft to identify the most important sources and references
- Extract information about studies, papers, authors, and findings mentioned
- Create proper APA format citations for each source
- Include author names, publication years, titles, and sources where possible
- If specific details are missing, create reasonable citations based on the content
- Limit the bibliography to a maximum of 10 references
- Choose the most relevant and important sources only
- Number the references if they are cited in the draft with [1], [2], etc.

Example format:
[1] Smith, J., & Johnson, A. (2023). Title of the study. Journal Name, 45(2), 123-145.

Bibliography (APA format, maximum 10 references):
"""

def get_bibliography_chain():
    """
    Returns a LangChain LLMChain for bibliography generation.
    """
    prompt = PromptTemplate(
        input_variables=["research_notes", "draft"],
        template=BIBLIOGRAPHY_PROMPT
    )
    
    # Always enable streaming in the LLM - callbacks will be passed during execution
    llm = OpenAI(
        temperature=0.1,
        openai_api_key=os.getenv("OPENAI_API_KEY"),
        streaming=True  # Always enable streaming
    )
    
    return LLMChain(llm=llm, prompt=prompt, output_key="bibliography")