from langchain_groq import ChatGroq
from langgraph.graph import StateGraph, START, END
from typing import TypedDict

GROQ_API_KEY = "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"

model = ChatGroq(
    api_key=GROQ_API_KEY,
    model="llama-3.3-70b-versatile",
    temperature=0.7,
)

class GraphState(TypedDict):
    input: str
    output: str


def chatbot(state: GraphState) -> dict:
  
    response = model.invoke(state["input"])
    return {"output": response.content}


builder = StateGraph(GraphState)
builder.add_node("chatbot", chatbot)
builder.add_edge(START, "chatbot")
builder.add_edge("chatbot", END)
graph = builder.compile()

def main():
    print("Simple AI Agent")
    print("Type 'exit' to quit.\n")

    while True:
        try:
            message = input("You: ")
        except (EOFError, KeyboardInterrupt):
            print()
            break

        if message.strip().lower() == "exit":
            break

        result = graph.invoke({"input": message})

        print(f"\nAI: {result['output']}\n")

if __name__ == "__main__":
    main()