Spaces:
Sleeping
Sleeping
| 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("<div style='text-align:center;'> <h1> Strategic Debate Assistant </h1> </div>") | |
| 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() |