from dataclasses import dataclass from langgraph.graph import START, StateGraph, END from typing import TypedDict from agents import general_agent, excel_supervisor, video_supervisor import os from typing import List OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") os.environ["OPENAI_API_KEY"] = str(OPENAI_API_KEY) @dataclass class Question: task_id: str question: str Level: str file_name: str local_file_path: str|None = None def get_file_type(file_path: str) -> str: """Determine file type from extension.""" if not file_path: return "none" file_path = file_path.lower() if file_path.endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')): return "image" elif file_path.endswith(('.xlsx', '.xls', '.csv')): return "excel" elif file_path.endswith('.py'): return "python" elif file_path.endswith(('.mp3', '.wav', '.m4a', '.ogg')): return "audio" else: return "unknown" def answer_qery(question: str, thread_id: str = "default") -> str: """Ask the agent a question.""" config = {"configurable": {"thread_id": thread_id}, "recursion_limit": 50} try: result = video_supervisor.invoke({ "messages": [ {"role": "user", "content": question} ] }) return result["messages"][-1].content except Exception as e: return f"Error: {str(e)}" def ask_question(question: str, thread_id: str = "default") -> str: """Ask the agent a question.""" config = {"configurable": {"thread_id": thread_id}, "recursion_limit": 100} try: response = general_agent.invoke( {"messages": [{"role": "user", "content": question}]}, config=config ) return response["messages"][-1].content except Exception as e: return f"Error: {str(e)}" def ask_question_with_file(question: Question, thread_id: str = "default") -> str: """Ask the agent a question, with optional file analysis.""" q = question.question root_file = "./files" file_path = root_file + "/" + question.file_name if not question.file_name: return ask_question(q, thread_id) file_type = get_file_type(file_path) # Create enhanced question with file guidance if file_type == "image": enhanced_question = f"{q}\n\nThere is an image file at '{file_path}'. Use the analyze_image tool to examine it." elif file_type == "excel": enhanced_question = f"{q}\n\nFile path: {file_path}" result = excel_supervisor.invoke({ "messages": [ {"role": "user", "content": enhanced_question} ] }) return result["messages"][-1].content elif file_type == "python": enhanced_question = f"{q}\n\nThere is a Python file at '{file_path}'. Use the read_python_file tool to examine it." elif file_type == "audio": enhanced_question = f"{q}\n\nThere is an audio file at '{file_path}'. Use the transcribe_audio tool to process it." else: enhanced_question = f"{q}\n\nThere is a file at '{file_path}' but I'm not sure what type it is." return ask_question(enhanced_question, thread_id) def ask_question_youtube(question: Question) -> str: """Ask the agent a question, with optional file analysis.""" q = question.question result = video_supervisor.invoke({ "messages": [ {"role": "user", "content": q} ] }) return result["messages"][-1].content # State class State(TypedDict): question: Question decision: str answer: str # NODE FUNCTIONS - These are the ones that work with LangGraph def ask_question_node(state: State) -> dict: """Node function for questions without files.""" question_obj = state["question"] thread_id = f"test_{question_obj.task_id}" # Call your existing function answer = answer_qery(question_obj.question, thread_id) # Return dict to update state return {"answer": answer} def ask_question_with_file_node(state: State) -> dict: """Node function for questions with files.""" question_obj = state["question"] thread_id = f"test_{question_obj.task_id}" # Call your existing function answer = ask_question_with_file(question_obj, thread_id) # Return dict to update state return {"answer": answer} def ask_question_youtube_node(state: State) -> dict: """Node function for questions with files.""" question_obj = state["question"] # Call your existing function answer = ask_question_youtube(question_obj) # Return dict to update state return {"answer": answer} def router_node(state: State): """Router node - returns dict to update state""" if state["question"].file_name: decision = "query_with_file" elif "youtube.com" in state["question"].question or "youtu.be" in state["question"].question: decision = "youtube" else: decision = "query" return {"decision": decision} def router_function(state: State): """Routing function - returns string to choose path""" return state["decision"] def build_graph(): # Graph builder = StateGraph(State) # Use the NODE functions (not the original functions) builder.add_node("query_with_file", ask_question_with_file_node) builder.add_node("query", ask_question_node) builder.add_node("youtube", ask_question_youtube_node) builder.add_node("router", router_node) # Define edges builder.add_edge(START, "router") builder.add_conditional_edges( "router", router_function, { "query_with_file": "query_with_file", "query": "query", "youtube": "youtube", }, ) builder.add_edge("query_with_file", END) builder.add_edge("query", END) builder.add_edge("youtube", END) react_graph = builder.compile() return react_graph def extract_final_answer(text: str) -> str|List[str]: """Extract the final answer from a string containing 'FINAL ANSWER: answer'""" # Method 1: Simple string split (most common case) if "FINAL ANSWER:" in text: # Split on "FINAL ANSWER:" and take the part after it parts = text.split("FINAL ANSWER:", 1) # Split only on first occurrence return parts else: return "FINAL ANSWER: unknown" if __name__ == "__main__": test = [ { "task_id": "8e867cd7-cff9-4e6c-867a-ff5ddc2550be", "question": "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia.", "Level": "1", "file_name": "" }, { "task_id": "99c9cc74-fdc8-46c6-8f8d-3ce2d3bfeea3", "question": "Hi, I'm making a pie but I could use some help with my shopping list. I have everything I need for the crust, but I'm not sure about the filling. I got the recipe from my friend Aditi, but she left it as a voice memo and the speaker on my phone is buzzing so I can't quite make out what she's saying. Could you please listen to the recipe and list all of the ingredients that my friend described? I only want the ingredients for the filling, as I have everything I need to make my favorite pie crust. I've attached the recipe as Strawberry pie.mp3.\n\nIn your response, please only list the ingredients, not any measurements. So if the recipe calls for \"a pinch of salt\" or \"two cups of ripe strawberries\" the ingredients on the list would be \"salt\" and \"ripe strawberries\".\n\nPlease format your response as a comma separated list of ingredients. Also, please alphabetize the ingredients.", "Level": "1", "file_name": "99c9cc74-fdc8-46c6-8f8d-3ce2d3bfeea3.mp3" } ] questions = [Question(**item) for item in test] for i, question in enumerate(questions): print(f"\n{i}. {question.question}") react_graph = build_graph() # Invoke the graph and capture the result result = react_graph.invoke({ "question": question, "decision": "", "answer": "" }) answer = result['answer'] print(answer) answer = extract_final_answer(answer)[1] print(answer)