The generate prompts function in the code should make three separate LLM calls with the same config and pass these prompts to the evaluator to select the best one.
# Import necessary libraries
import os
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langgraph.graph import StateGraph, END
from typing import List, Dict, Any, TypedDict
import json
os.environ["OPENAI_API_KEY"] = "<TRUFFLEHOG_REDACTED_OPENAI>"
# Define TypedDicts for structured state management
class OutputField(TypedDict):
    name: str
    description: str
    type: str
    source: str
    required: bool

class WorkflowState(TypedDict):
    goal: str
    output_fields: List[OutputField]
    example_queries: List[Dict[str, str]]
    prompts: Dict[str, str]
    evaluations: Dict[str, Any]
    best_prompt: Dict[str, Any]

# Define Agents

def enhance_text(state: WorkflowState, field_name: str, content: str) -> WorkflowState:
    """
    Enhances a given text (Goal or Output Field) using an LLM.
    """
    llm = ChatOpenAI(model="gpt-4o-mini-2024-07-18", temperature=0)
    prompt = PromptTemplate.from_template("""
    You are an expert in refining and improving text.

    Original {field}: {content}

    Improve the above {field} for better clarity and effectiveness.

    Improved {field}:
    """)
    chain = LLMChain(prompt=prompt, llm=llm)
    improved_response = chain.run(field=field_name, content=content)
    
    # Access the content of the response
    if isinstance(improved_response, AIMessage):
        improved_text = improved_response.content.strip()
    else:
        improved_text = improved_response.strip()
    
    # Update the appropriate field in the state
    if field_name.lower() == "goal":
        state["goal"] = improved_text
    elif field_name.startswith("output_field_"):
        # Extract the index from the field name (e.g., "output_field_1" -> 1)
        try:
            index = int(field_name.split("_")[-1]) - 1
            if 0 <= index < len(state["output_fields"]):
                state["output_fields"][index]["description"] = improved_text
        except (ValueError, IndexError):
            pass
    
    return state
def enhance_output_fields_simulated(state: WorkflowState) -> WorkflowState:
    """
    Simulates enhancement for specific Output Fields.
    """
    # Create a copy of output fields to avoid modifying during iteration
    output_fields = state["output_fields"].copy()
    
    # For testing, enhance each output field
    for idx, field in enumerate(output_fields, start=1):
        enhanced_state = enhance_text(
            state=state,
            field_name=f"output_field_{idx}",
            content=field["description"]
        )
        state = enhanced_state  # Update the state with enhanced descriptions
        print(f"Output Field '{field['name']}' has been enhanced.\n")
    
    return state

def generate_example_queries(state: WorkflowState) -> WorkflowState:
    """
    Generates three example queries based on the Goal and Output Fields.
    """
    llm = ChatOpenAI(model="gpt-4o-mini-2024-07-18", temperature=0)
    prompt = PromptTemplate.from_template("""
    Given the following Goal and Output Fields, generate three example queries that a user might ask.

    Goal: {goal}

    Output Fields:
    {output_fields}

    Example Queries:
    1.
    2.
    3.
    """)
    output_fields_str = "\n".join([
        f"- **Name**: {field['name']}\n  **Description**: {field['description']}\n  **Type**: {field['type']}\n  **Source**: {field['source']}\n  **Required**: {field['required']}"
        for field in state["output_fields"]
    ])
    chain = LLMChain(prompt=prompt, llm=llm)
    examples = chain.run(goal=state["goal"], output_fields=output_fields_str).strip()
    # Parse the generated examples into a list
    example_queries = []
    for line in examples.split('\n'):
        line = line.strip()
        if line.startswith('1.') or line.startswith('2.') or line.startswith('3.'):
            query = line.split('.', 1)[1].strip()
            if query:
                example_queries.append({"query": query, "response": ""})
    state["example_queries"] = example_queries
    return state

def generate_responses(state: WorkflowState) -> WorkflowState:
    """
    Generates responses for each example query based on the Goal and Output Fields.
    """
    llm = ChatOpenAI(model="gpt-4o-mini-2024-07-18", temperature=0)
    prompt = PromptTemplate.from_template("""
    You are an assistant that provides clear and concise responses to user queries based on the following Goal and Output Fields.

    Goal: {goal}

    Output Fields:
    {output_fields}

    Query: {query}

    Response:
    """)

    goal = state.get("goal", "")
    output_fields = state.get("output_fields", [])
    output_fields_str = "\n".join([
        f"- **Name**: {field['name']}\n  **Description**: {field['description']}\n  **Type**: {field['type']}\n  **Source**: {field['source']}\n  **Required**: {field['required']}"
        for field in output_fields
    ])

    for example in state.get("example_queries", []):
        query = example.get("query", "")
        if not query:
            example["response"] = "No query provided."
            continue

        # Prepare the prompt
        filled_prompt = prompt.format(
            goal=goal,
            output_fields=output_fields_str,
            query=query
        )

        # Generate response
        response = llm(filled_prompt).content.strip()
        example["response"] = response

    return state

def generate_prompts(state: WorkflowState) -> WorkflowState:
    """
    Generates three distinct prompts (A, B, C) using separate LLM calls and saves them in the state.
    """
    llm = ChatOpenAI(model="gpt-4o-mini-2024-07-18", temperature=0)
    prompt_template = PromptTemplate.from_template("""
    Using the following Goal, Output Fields, and Example Queries with Responses, create a detailed prompt for an AI assistant.
    Make each variation distinct in its approach.

    Goal: {goal}

    Output Fields:
    {output_fields}

    Example Queries and Responses:
    {examples}

    Create a {variation} style prompt that emphasizes:
    A: Clarity and precision in responses
    B: Comprehensive coverage of all fields
    C: Natural conversation flow and user engagement

    Generated Prompt:
    """)

    variations = ["A", "B", "C"]
    prompts = {}

    # Ensure we have the required state data
    if not state.get("example_queries"):
        state["example_queries"] = []

    examples_str = "\n".join([
        f"**Query**: {ex.get('query', '')}\n**Response**: {ex.get('response', '')}"
        for ex in state["example_queries"]
    ])

    output_fields_str = "\n".join([
        f"- **Name**: {field['name']}\n  **Description**: {field['description']}\n  **Type**: {field['type']}\n  **Source**: {field['source']}\n  **Required**: {field['required']}"
        for field in state["output_fields"]
    ])

    for var in variations:
        filled_prompt = prompt_template.format(
            goal=state.get("goal", ""),
            output_fields=output_fields_str,
            examples=examples_str,
            variation=var
        )
        generated_prompt = llm(filled_prompt).content.strip()  # Call LLM to generate each prompt
        prompts[f"prompt_{var.lower()}"] = generated_prompt  # Store generated prompt

    state["prompts"] = prompts  # Save the prompts in state
    return state
# Add error handling to evaluate_prompts
def evaluate_prompts(state: WorkflowState) -> WorkflowState:
    """
    Evaluates the three generated prompts and selects the best one.
    """
    llm = ChatOpenAI(model="gpt-4o-mini-2024-07-18", temperature=0)
    prompt = PromptTemplate.from_template("""
    You are G-Eval, an evaluator that assesses the completeness and clarity of AI prompts.
    Rate each prompt on:
    1. Completeness (covers all requirements)
    2. Clarity (easy to understand)
    3. Practicality (feasible to implement)

    Evaluate the following prompts and rate each on a scale of 1 to 10:

    Prompt A:
    {prompt_a}

    Prompt B:
    {prompt_b}

    Prompt C:
    {prompt_c}

    Return a JSON object with this exact format:
    {
        "prompt_a": <score>,
        "prompt_b": <score>,
        "prompt_c": <score>,
        "selected_prompt": "<winning_prompt_letter>"
    }
    """)
    
    # Ensure prompts exist and are non-empty
    prompt_a = state["prompts"].get("prompt_a", "").strip()
    prompt_b = state["prompts"].get("prompt_b", "").strip()
    prompt_c = state["prompts"].get("prompt_c", "").strip()
    
    if not prompt_a or not prompt_b or not prompt_c:
        state["evaluations"] = {
            "prompt_a": 0,
            "prompt_b": 0,
            "prompt_c": 0,
            "selected_prompt": "prompt_a"
        }
        return state
    
    chain = LLMChain(prompt=prompt, llm=llm)
    
    try:
        evaluation = chain.run(
            prompt_a=prompt_a,
            prompt_b=prompt_b,
            prompt_c=prompt_c
        ).strip()
        
        evaluation_dict = json.loads(evaluation)
        
        # Validate the evaluation format
        required_keys = ["prompt_a", "prompt_b", "prompt_c", "selected_prompt"]
        if not all(key in evaluation_dict for key in required_keys):
            raise ValueError("Missing required keys in evaluation response")
            
    except (json.JSONDecodeError, ValueError, Exception) as e:
        print(f"Error in evaluation: {str(e)}. Response: {evaluation if 'evaluation' in locals() else 'No response'}")
        evaluation_dict = {
            "prompt_a": 0,
            "prompt_b": 0,
            "prompt_c": 0,
            "selected_prompt": "prompt_a"
        }
    
    state["evaluations"] = evaluation_dict
    return state
        

# Define the Workflow

def collect_user_inputs_simulated(state: WorkflowState) -> WorkflowState:
    """
    Simulates user inputs for testing purposes.
    """
    # Placeholder example for recipe search
    goal = "Find recipes that require overnight marinating."

    output_fields = [
        {
            "name": "Recipe Name",
            "description": "The name of the recipe.",
            "type": "String",
            "source": "Recipe Database",
            "required": True
        },
        {
            "name": "Preparation Time",
            "description": "Total time required to prepare the recipe, including marinating time.",
            "type": "Integer",
            "source": "Recipe Database",
            "required": True
        },
        {
            "name": "Ingredients",
            "description": "List of ingredients needed for the recipe.",
            "type": "List",
            "source": "Recipe Database",
            "required": True
        }
    ]

    # Update the state with these values
    state["goal"] = goal
    state["output_fields"] = output_fields
    state["example_queries"] = []
    state["prompts"] = {}
    state["evaluations"] = {}
    state["best_prompt"] = {}

    print("Simulated user inputs collected.\n")
    return state

# Create and Compile the Workflow Graph

def build_workflow_graph() -> StateGraph:
    """
    Constructs and compiles the workflow graph.
    """
    workflow = StateGraph(WorkflowState)
    
    # Node: Collect User Inputs (Simulated for Testing)
    workflow.add_node("collect_inputs", collect_user_inputs_simulated)
    
    # Node: Enhance Goal
    workflow.add_node("enhance_goal", lambda state: enhance_text(state, field_name="Goal", content=state["goal"]))
    
    # Node: Enhance Output Fields
    workflow.add_node("enhance_output_fields", enhance_output_fields_simulated)
    
    # Node: Generate Example Queries
    workflow.add_node("generate_queries", generate_example_queries)
    
     # Node: Modify Example Queries (Simulated: Accept as is)
    def modify_example_queries_simulated(state: WorkflowState) -> WorkflowState:
        """
        Simulates user modification of example queries by accepting them as generated.
        """
        if state["example_queries"]:
            print("Modifying Example Queries: Accepting as generated.\n")
        else:
            print("No example queries to modify.\n")
        return state
    
    workflow.add_node("modify_queries", modify_example_queries_simulated)
    
    # Node: Modify Example Queries
    # Node: Generate Responses
    workflow.add_node("generate_responses", generate_responses)
    
    # Node: Generate Prompts
    workflow.add_node("generate_prompts", generate_prompts)
    
    # Node: Evaluate Prompts
    workflow.add_node("evaluate_prompts", evaluate_prompts)
    
    # Node: Select Best Prompt
    workflow.add_node("select_best_prompt", select_best_prompt)
    
    # Define the workflow transitions
    workflow.set_entry_point("collect_inputs")
    workflow.add_edge("collect_inputs", "enhance_goal")
    workflow.add_edge("enhance_goal", "enhance_output_fields")
    workflow.add_edge("enhance_output_fields", "generate_queries")
    workflow.add_edge("generate_queries", "modify_queries")
    workflow.add_edge("modify_queries", "generate_responses")
    workflow.add_edge("generate_responses", "generate_prompts")
    workflow.add_edge("generate_prompts", "evaluate_prompts")
    workflow.add_edge("evaluate_prompts", "select_best_prompt")
    workflow.add_edge("select_best_prompt", END)
    
    return workflow.compile()

# Main Execution Function

def run_workflow():
    """
    Executes the workflow and displays the best prompt.
    """
    # Ensure the OpenAI API key is set
    if "OPENAI_API_KEY" not in os.environ:
        api_key = input("Enter your OpenAI API Key: ").strip()
        if not api_key:
            print("OpenAI API Key is required to run this program.")
            return
        os.environ["OPENAI_API_KEY"] = api_key
    
    # Build the workflow graph
    workflow = build_workflow_graph()
    
    # Initialize the state
    initial_state: WorkflowState = {
        "goal": "",
        "output_fields": [],
        "example_queries": [],
        "prompts": {},
        "evaluations": {},
        "best_prompt": {}
    }
    
    # Invoke the workflow
    print("Running the workflow...\n")
    final_state = workflow.invoke(initial_state)
    
    # Display the best prompt
    selected = final_state.get("best_prompt", {}).get("selected_prompt", "A")
    prompt = final_state.get("best_prompt", {}).get("prompt", "No prompt generated.")
    
    print("\n=== Best Prompt Selected ===")
    print(f"Prompt {selected}:\n")
    print(prompt)

# Run the workflow
run_workflow()