Spaces:
Running
Running
| from typing import TypedDict | |
| from tavily import TavilyClient | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| from langgraph.graph import StateGraph, START, END | |
| import gradio as gr | |
| from langgraph.checkpoint.memory import InMemorySaver | |
| import os | |
| class GraphState(TypedDict): | |
| query: str | |
| search_results1: str | |
| search_results2: str | |
| summary1: str | |
| summary2: str | |
| tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY")) | |
| def tavily_search(state: GraphState): | |
| response = tavily_client.search( | |
| query=state["query"], | |
| max_results=2 | |
| ) | |
| return { | |
| "search_results1": response["results"][0]['content'][:500], | |
| "search_results2": response["results"][1]['content'][:500] | |
| } | |
| llm = ChatGoogleGenerativeAI( | |
| model="gemini-2.5-flash", | |
| google_api_key=os.getenv("GOOGLE_API_KEY"), | |
| ) | |
| def summarize1(state: GraphState): | |
| prompt = f""" | |
| Simply summarize following: | |
| {state['search_results1']} | |
| """ | |
| response = llm.invoke(prompt) | |
| return {"summary1": response.content} | |
| def summarize2(state: GraphState): | |
| prompt = f""" | |
| Simply summarize following: | |
| {state['search_results2']} | |
| """ | |
| response = llm.invoke(prompt) | |
| return {"summary2": response.content} | |
| builder = StateGraph(GraphState) | |
| builder.add_node("search", tavily_search) | |
| builder.add_node("summarize", summarize1) | |
| builder.add_node("revise", summarize2) | |
| builder.set_entry_point("search") | |
| builder.add_edge("search", "summarize") | |
| builder.add_edge("summarize", "revise") | |
| builder.add_edge("revise", END) | |
| checkpointer = InMemorySaver() | |
| thread = {"configurable": {"thread_id": "1"}} | |
| graph = builder.compile( | |
| checkpointer=checkpointer, | |
| interrupt_after=["summarize"] | |
| ) | |
| def describe_topic(topic): | |
| result = graph.invoke( | |
| {"query": topic}, | |
| config=thread | |
| ) | |
| print(result["search_results1"]) | |
| return result["summary1"] | |
| def notgood(topic): | |
| result = graph.invoke( | |
| None, | |
| config=thread | |
| ) | |
| print(result["search_results2"]) | |
| return result["summary2"] | |
| with gr.Blocks() as demo: | |
| gr.Markdown("""# Ask anything and press Enter""") | |
| topic_input = gr.Textbox( | |
| label="Ask here" | |
| ) | |
| output1 = gr.Textbox( | |
| label="AI description" | |
| ) | |
| output2 = gr.Textbox( | |
| label="AI revised description" | |
| ) | |
| no_btn = gr.Button("Click here to revise description") | |
| topic_input.submit( | |
| describe_topic, | |
| inputs=topic_input, | |
| outputs=output1 | |
| ) | |
| no_btn.click( | |
| notgood, | |
| inputs=topic_input, | |
| outputs=output2 | |
| ) | |
| gr.Markdown(""" | |
| **Note:** | |
| I'm using Google Gemini **free API tier**. | |
| I sacrificed some **quality** by limiting **LLM input output** to not exceed free tier. | |
| If "error" shows up, it means free tier exceeded. | |
| """) | |
| gr.Image( | |
| value="graph.png" | |
| ) | |
| demo.launch() |