app1.py

import streamlit as st
from typing import TypedDict, Optional
from langgraph.graph import StateGraph, END
from together import Together
import os
import re
from langchain.agents import initialize_agent, AgentType
from tools import get_all_tools
from langchain_together import ChatTogether

# --------------------------------------------
# LangChain Tool Agent Setup (Router + Tools)
# --------------------------------------------

tools = get_all_tools()
llm = ChatTogether(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    together_api_key=os.environ["together_api_key"]
)
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

def decide_route_from_output(output: str) -> str:
    output = output.lower()
    if "explain" in output:
        return "explain_code"
    elif "generate" in output or "code:" in output:
        return "code_generator"
    elif "execution result" in output or "run this code" in output:
        return "execute_code"
    elif "thought" in output or "reason" in output:
        return "deep_think"
    else:
        return "END"

def router_node(state: dict) -> dict:
    user_input = state.get("input", "")
    result = agent.run(user_input)
    route = decide_route_from_output(result)
    return {**state, "output": result, "route": route}

# Initialize Together client (for raw completions)
client = Together(api_key="d03eaa038011423d1b8828a3ff2956511b2ca34b6cc88ffcbc4ad6314db31319")

def llm_invoke(prompt: str) -> str:
    try:
        response = client.completions.create(
            model="mistralai/Mixtral-8x7B-Instruct-v0.1",
            prompt=prompt,
            max_tokens=3000,
            temperature=0.2,
            stop=["</s>"]
        )
        return response.choices[0].text.strip()
    except Exception as e:
        st.error(f"LLM API Error: {str(e)}")
        return "Error in LLM response"

# ----------------------
# State Definition
# ----------------------
class AgentState(TypedDict):
    domain: str
    goal: str
    follow_up_questions: Optional[str]
    follow_up_answers: Optional[str]
    agents_needed: Optional[str]
    code_generator: Optional[str]
    explanation: Optional[str]
    input: Optional[str]
    output: Optional[str]
    route: Optional[str]

# ----------------------
# Agent Functions
# ----------------------
def collect_domain_goal(state: AgentState) -> AgentState:
    return state

def needs_follow_up(goal: str) -> bool:
    prompt = f"""
    Does this system goal require clarifying questions about input data, file structure, or formats?

    Goal: {goal}

    Respond only 'yes' or 'no'.
    """
    response = llm_invoke(prompt).strip().lower()
    return response == "yes"

def ask_follow_up_questions(state: AgentState) -> AgentState:
    if needs_follow_up(state["goal"]):
        prompt = f"""
        Based on the domain: {state['domain']} and goal: {state['goal']},
        list 2–3 critical follow-up questions about:
        - Data/file formats (CSV, JSON, etc.)
        - Column names or structure
        - Required reference files (terms, conditions, etc.)
        - Any APIs or databases used

        Return in bullet-point format.
        """
        state["follow_up_questions"] = llm_invoke(prompt)
    else:
        state["follow_up_questions"] = "No questions needed"
    return state

def identify_agents(state: AgentState) -> AgentState:
    prompt = f"""
    Domain: {state['domain']}
    Goal: {state['goal']}
    Context: {state.get('follow_up_answers', 'No additional info')}

    List the agents needed in an agentic system. For each:
    - Name
    - Purpose
    - What it processes or validates

    Return as a clean markdown list.
    """
    state["agents_needed"] = llm_invoke(prompt)
    return state

def generate_code(state: AgentState) -> AgentState:
    prompt = f"""
    You are an expert Python engineer. Based on:
    - Domain: {state['domain']}
    - Goal: {state['goal']}
    - Follow-up Answers: {state.get('follow_up_answers', 'None')}
    - Agents Needed: {state.get('agents_needed', 'None')}

    Generate a full working LangGraph app using this structure:
    1. Imports: streamlit, langgraph, together, re, os, csv if needed
    2. State: TypedDict
    3. One clean agent function per node
    4. LangGraph construction with node/edge logic
    5. Together LLM call helper
    6. Streamlit UI to collect user input and run graph
    7. Output final result cleanly

    Match code style and flow of our health advisory example. Avoid AgentState subclasses. Ensure the system is executable.
    """
    state["code_generator"] = llm_invoke(prompt)
    return state

def explain_code(state: AgentState) -> AgentState:
    if state.get("code_generator"):
        prompt = f"Explain this code simply and clearly:\n\n{state['code_generator']}"
        state["explanation"] = llm_invoke(prompt)
    else:
        state["explanation"] = "No code to explain"
    return state

# ----------------------
# Graph Construction
# ----------------------
def create_graph():
    graph = StateGraph(AgentState)
    graph.add_node("collect_domain_goal", collect_domain_goal)
    graph.add_node("ask_follow_up_questions", ask_follow_up_questions)
    graph.add_node("identify_agents", identify_agents)
    graph.add_node("generate_code", generate_code)
    graph.add_node("explain_code", explain_code)

    graph.set_entry_point("collect_domain_goal")
    graph.add_edge("collect_domain_goal", "ask_follow_up_questions")
    graph.add_edge("ask_follow_up_questions", "identify_agents")
    graph.add_edge("identify_agents", "generate_code")
    graph.add_edge("generate_code", "explain_code")
    graph.add_edge("explain_code", END)

    return graph.compile()

app = create_graph()

# ----------------------
# Streamlit Interface
# ----------------------
def main():
    st.set_page_config(page_title="Agentic Code Generator", layout="centered")
    st.title("🤖 Auto-Build LangGraph App")

    domain = st.text_input("Domain (e.g. finance, warranty claims, medical):")
    goal = st.text_area("System Goal (be detailed, include what data is needed):", height=150)

    if "state" not in st.session_state:
        st.session_state.state = {
            "domain": "",
            "goal": "",
            "follow_up_questions": None,
            "follow_up_answers": None,
            "agents_needed": None,
            "code_generator": None,
            "explanation": None,
            "input": None,
            "output": None,
            "route": None
        }

    if "step" not in st.session_state:
        st.session_state.step = "initial"

    if st.button("🚀 Generate Agentic App") and domain and goal:
        st.session_state.state.update({"domain": domain, "goal": goal})
        st.session_state.step = "processing"

        with st.spinner("Thinking like an architect..."):
            try:
                result = app.invoke(st.session_state.state)
                st.session_state.state.update(result)

                if (st.session_state.state.get("follow_up_questions") and
                    st.session_state.state["follow_up_questions"] != "No questions needed"):
                    st.session_state.step = "follow_up"
                else:
                    st.session_state.step = "complete"

            except Exception as e:
                st.error(f"Error processing request: {str(e)}")
                st.session_state.step = "error"

    if st.session_state.step == "follow_up":
        st.subheader("🔍 Follow-Up Questions")
        st.markdown(st.session_state.state["follow_up_questions"])

        answers = st.text_area("Your answers to the above:", height=150)
        if st.button("Continue with those answers"):
            st.session_state.state["follow_up_answers"] = answers
            st.session_state.step = "processing_final"
            with st.spinner("Building your LangGraph app..."):
                try:
                    final_result = app.invoke(st.session_state.state)
                    st.session_state.state.update(final_result)
                    st.session_state.step = "complete"
                except Exception as e:
                    st.error(f"Error in final processing: {str(e)}")
                    st.session_state.step = "error"

    if st.session_state.step == "complete":
        st.subheader("🎯 Identified Agents")
        st.markdown(st.session_state.state["agents_needed"])

        st.subheader("✅ Generated Code")
        st.code(st.session_state.state["code_generator"], language="python")

        st.subheader("🧠 Code Explanation")
        st.markdown(st.session_state.state["explanation"])

        st.download_button(
            label="📥 Download Generated Code",
            data=st.session_state.state["code_generator"],
            file_name="generated_langgraph_app.py",
            mime="text/python"
        )

        st.subheader("🧰 Use Agentic Tools")
        _ = st.text_input("Optional prompt for tools:", key="tool_input")

        col1, col2, col3 = st.columns(3)
        with col1:
            if st.button("Explain Code"):
                st.session_state.state["input"] = "Explain this code:\n" + st.session_state.state["code_generator"]
                st.session_state.step = "router_tool"
        with col2:
            if st.button("Run Code"):
                st.session_state.state["input"] = "Run this code:\n" + st.session_state.state["code_generator"]
                st.session_state.step = "router_tool"
        with col3:
            if st.button("Deep Think"):
                st.session_state.state["input"] = "Think deeply about this:\n" + st.session_state.state["goal"]
                st.session_state.step = "router_tool"

    if st.session_state.step == "router_tool":
        with st.spinner("Running selected tool..."):
            try:
                result = router_node(st.session_state.state)
                st.session_state.state.update(result)
                st.session_state.step = "tool_complete"
            except Exception as e:
                st.error(f"Tool execution failed: {e}")
                st.session_state.step = "error"

    if st.session_state.step == "tool_complete":
        st.subheader("🔍 Tool Output")
        st.markdown(st.session_state.state.get("output", "No output available."))

    if st.button("🔄 Reset"):
        st.session_state.clear()
        st.experimental_rerun()

if __name__ == "__main__":
    main()