import token from gradio import Interface, Textbox, Markdown, Slider, Chatbot, ChatMessage, Blocks, Row, Column, ClearButton, Button, Number,Label from final_graph import create_final_graph, render_bettafish_report from forum_graph import create_forum_graph from graph_states import FinalState, ForumState from langchain_core.callbacks import UsageMetadataCallbackHandler import uuid import time graph = None async def init_graph(): global graph if graph is None: graph = await create_final_graph() async def process_inputs(query, refinements, max_rounds): state = FinalState( query=query, vector_store="", query_limit=int(refinements), max_rounds=int(max_rounds), current_round=1, step=0, round_digests=[], running_mermaid_graph="", debate_history=[], final_report=None, messages = [], references = [] ) await init_graph() messages = [] thread_id = str(uuid.uuid4()) usage_callback = UsageMetadataCallbackHandler() config = {"configurable":{"thread_id":thread_id},"recursion_limit":100,"callbacks":[usage_callback]} step = 0 current_round = 1 current_round_md = "### Current Round: {current_round}" start_time = time.time() token_used = 0 input_tokens = 0 output_tokens = 0 role = "user" status = "Searching the web for relevant information..." try: async for parent,child in graph.astream(state,config=config,stream_mode="updates",subgraphs=True): print("Parent:",parent) print("Child:",child,end="\n\n") if not parent or len(parent)==0: continue parent_node = parent[0].split(":")[0] curr_node = list(child.keys())[0] values = child.get(curr_node,{}) if parent_node == "Research": if curr_node == "Query Refiner": status = "Searching with refined queries..." elif parent_node == "Forum Graph": if curr_node == "Debate Supervisor": step = values.get("step",step) msgs = values.get("messages",[]) if step%2 == 0: role = "assistant" else: role = "user" if len(msgs)>0: messages.append(ChatMessage(role=role,content=msgs[-1].content,metadata={"title":"Moderator Instructions"})) status = "Analyzing and creating arguments..." elif curr_node == "Persona Agent": msgs = values.get("messages",[]) if msgs: tool_call = False curr_message = msgs[-1] if len(curr_message.content.strip()) == 0 and len(curr_message.tool_calls)>0: tool_call = True if not tool_call: messages.append(ChatMessage(role=role,content=curr_message.content)) else: messages.append(ChatMessage(role=role,content="Tool Call Made. Awaiting Response...",metadata={"title":"Fetching researched data"})) status = "Engaging in debate..." elif curr_node == "Round Digest": current_round += 1 status = "Preparing for next round..." for model, usage_data in usage_callback.usage_metadata.items(): input_tokens += usage_data.get("input_tokens",0) output_tokens += usage_data.get("output_tokens",0) token_used += usage_data.get("total_tokens",0) total_time_taken = round((time.time()-start_time),2) yield messages, "" , status, current_round_md.format(current_round=current_round),total_time_taken,input_tokens,output_tokens, token_used except Exception as e: print("Error during running the graph:",e) yield messages, "" , f" An error occurred during the graph execution: {e} ", current_round_md.format(current_round=current_round),total_time_taken,input_tokens,output_tokens, token_used return try: snapshot = graph.get_state(config={"configurable":{"thread_id":thread_id}}) if snapshot.values: report=render_bettafish_report(snapshot.values) else: report = "No state found. Did the graph finish?" yield messages, report , " Graph Execution Completed! Report Generated Below. ", current_round_md.format(current_round=current_round),total_time_taken,input_tokens,output_tokens, token_used except Exception as e: print("Error during report generation:",e) yield messages, "" , f" An error occurred during report generation: {e} ", current_round_md.format(current_round=current_round),total_time_taken,input_tokens,output_tokens, token_used return with Blocks() as demo: Markdown("

Strategic Debate Assistant

") with Row(): with Column(): with Row(): refinements = Slider(0, 3, label="Max Query Refinements", step=1, value=0) max_rounds = Slider(1, 10, label="Max Debate Rounds",step=1, value=1) query = Textbox(label="Enter your strategic question here", lines=2) with Row(): ClearButton([query,refinements,max_rounds], value="Reset") submit_btn = Button("Submit",variant="primary") with Row(): input_tokens = Number(label=" Input Tokens Used", value=0,interactive=False) output_tokens = Number(label="Output Tokens Used", value=0,interactive=False) total_tokens = Number(label="Total Tokens Used", value=0,interactive=False) total_time = Number(label="Total Time Taken (s)", value=0,interactive=False) label = Textbox(value=" Report will be generated below upon submission", label="Status", interactive=False ) with Column(): current_round = Markdown(label="Current round") chatbot = Chatbot(label="Debate Conversation") report_md = Markdown(label="Strategic Report") submit_btn.click(process_inputs, inputs=[query, refinements, max_rounds], outputs=[chatbot, report_md,label, current_round, total_time,input_tokens,output_tokens,total_tokens]) demo.launch()