Spaces:
Sleeping
Sleeping
| from typing import Annotated, List, Dict | |
| from typing_extensions import TypedDict | |
| from langgraph.graph import END, START, StateGraph | |
| from langgraph.graph.message import add_messages | |
| from langgraph.prebuilt import ToolNode, tools_condition | |
| from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, BaseMessage | |
| import json | |
| from config import llm, logger | |
| from tools import fitness_analysis_tool | |
| class State(TypedDict): | |
| messages : Annotated[list, add_messages] | |
| youtube_url : str | |
| user_video : str | |
| groq_api_key : str | |
| result : dict | |
| SYSTEM_PROMPT = """ | |
| You are an AI fitness coach. You have one tool: fitness_analysis_tool. | |
| When given a YouTube URL and a user video path, IMMEDIATELY call | |
| fitness_analysis_tool with: | |
| - youtube_url = the YouTube URL provided | |
| - user_video_path = the user video path provided | |
| - groq_api_key = the groq api key provided | |
| Do NOT ask questions. Do NOT ask for context. | |
| Just call the tool immediately with the provided parameters. | |
| """ | |
| def make_tool_graph(): | |
| tools = [fitness_analysis_tool] | |
| tool_node = ToolNode(tools) | |
| llm_with_tools = llm.bind_tools(tools) | |
| def call_llm(state: State): | |
| messages = state["messages"] | |
| # ── Normalize: accept both dicts and LangChain message objects ── | |
| normalized = [] | |
| for m in messages: | |
| if isinstance(m, BaseMessage): | |
| # Already a LangChain message object — use as-is | |
| normalized.append(m) | |
| elif isinstance(m, dict): | |
| role = m.get("role", "user") | |
| content = m.get("content", "") | |
| if role == "system": | |
| normalized.append(SystemMessage(content=content)) | |
| elif role == "assistant": | |
| normalized.append(AIMessage(content=content)) | |
| else: | |
| normalized.append(HumanMessage(content=content)) | |
| else: | |
| # Fallback — convert to string | |
| normalized.append(HumanMessage(content=str(m))) | |
| # ── Prepend system prompt if not present ── | |
| has_system = any(isinstance(m, SystemMessage) for m in normalized) | |
| if not has_system: | |
| normalized = [SystemMessage(content=SYSTEM_PROMPT)] + normalized | |
| # ── Inject context into last HumanMessage ── | |
| for i in range(len(normalized) - 1, -1, -1): | |
| if isinstance(normalized[i], HumanMessage): | |
| ctx = "" | |
| if state.get("youtube_url"): | |
| ctx += f"\nYouTube URL: {state['youtube_url']}" | |
| if state.get("user_video"): | |
| ctx += f"\nUser video path: {state['user_video']}" | |
| if state.get("groq_api_key"): | |
| ctx += f"\nGroq API key: {state['groq_api_key']}" | |
| if ctx: | |
| normalized[i] = HumanMessage(content=normalized[i].content + ctx) | |
| break | |
| try: | |
| response = llm_with_tools.invoke(normalized) | |
| return {"messages": [response]} | |
| except Exception as e: | |
| logger.error(f"LLM call error: {e}") | |
| return { | |
| "messages": [AIMessage(content=f"Error: {str(e)}")] | |
| } | |
| builder = StateGraph(State) | |
| builder.add_node("llm", call_llm) | |
| builder.add_node("tools", tool_node) | |
| builder.add_edge(START, "llm") | |
| builder.add_conditional_edges("llm", tools_condition) | |
| builder.add_edge("tools", "llm") | |
| return builder.compile() | |
| tool_agent = make_tool_graph() |