utils.py

import json
from config import llm
from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages

class State(TypedDict):
    messages: Annotated[list, add_messages]

def format_make_infra_response(tool_result):
    """Format the structured output from make_infra tool for better readability"""
    try:
        # Try to parse if it's a JSON string
        if isinstance(tool_result, str):
            try:
                data = json.loads(tool_result)
            except:
                # If it's not JSON, return as is
                return tool_result
        else:
            data = tool_result
        
        if isinstance(data, dict):
            if data.get('success'):
                links_info = ""
                if data.get('build_url'):
                    # Job is running, provide direct build link
                    links_info = f"*🔗 Build Link:* <{data.get('build_url')}|View Build #{data.get('build_number')}>\n"
                else:
                    # Job not started yet, provide job page and queue links
                    links_info = f"*🔗 Job Page:* <{data.get('job_url')}|View All Builds>\n*📋 Queue:* <{data.get('queue_url')}|Queue Item #{data.get('queue_id')}>\n"
                
                return f"""✅ *Infrastructure Build Triggered Successfully*

*Job ID:* `{data.get('job_id')}`
*User:* {data.get('username')}
*Job Name:* {data.get('job_name')}
{links_info}
{data.get('message', '')}"""
            else:
                return f"""❌ *Infrastructure Build Failed*

*Error:* {data.get('error')}
*Job ID:* {data.get('job_id') or 'N/A'}"""
        else:
            return str(tool_result)
    except Exception as e:
        print(f"Error formatting make_infra response: {e}")
        return str(tool_result)


def format_update_image_digest_response(tool_result):
    """Format the structured output from update_image_digest tool for better readability"""
    try:
        # Try to parse if it's a JSON string
        if isinstance(tool_result, str):
            try:
                data = json.loads(tool_result)
            except:
                # If it's not JSON, return as is
                return tool_result
        else:
            data = tool_result
        
        if isinstance(data, dict):
            if data.get('success'):
                return f"""✅ *Image Digest Updated Successfully*

*Pull Request:* <{data.get('pr_url')}|PR #{data.get('pr_number')}>
*Branch:* `{data.get('branch_name')}`
*Config File:* `{data.get('config_file')}`

*Changes:*
• Old digest: `{data.get('old_digest')}`
• New digest: `{data.get('new_digest')}`

{data.get('message', '')}"""
            else:
                return f"""❌ *Image Digest Update Failed*

*Error:* {data.get('error')}
*PR URL:* {data.get('pr_url') or 'N/A'}"""
        else:
            return str(tool_result)
    except Exception as e:
        print(f"Error formatting update_image_digest response: {e}")
        return str(tool_result)


def enhanced_chatbot(state: State):
    """Enhanced chatbot that formats tool responses appropriately"""
    messages = state["messages"]
    
    # Check only the last message to see if it's a tool result
    if messages:
        last_msg = messages[-1]
        if hasattr(last_msg, 'name') and hasattr(last_msg, 'content'):
            if last_msg.name == "make_infra":
                try:
                    # Try to parse the JSON result and format it
                    data = json.loads(last_msg.content)
                    make_infra_result = format_make_infra_response(last_msg.content)
                    # Return the formatted result directly
                    from langchain_core.messages import AIMessage
                    return {"messages": [AIMessage(content=make_infra_result)]}
                except:
                    # If parsing fails, return the raw content
                    from langchain_core.messages import AIMessage
                    return {"messages": [AIMessage(content=last_msg.content)]}
            elif last_msg.name == "update_image_digest":
                try:
                    # Try to parse the JSON result and format it
                    data = json.loads(last_msg.content)
                    update_result = format_update_image_digest_response(last_msg.content)
                    # Return the formatted result directly
                    from langchain_core.messages import AIMessage
                    return {"messages": [AIMessage(content=update_result)]}
                except:
                    # If parsing fails, return the raw content
                    from langchain_core.messages import AIMessage
                    return {"messages": [AIMessage(content=last_msg.content)]}
    
    # If no tool result, generate normal response
    response = llm.llm_with_tools.invoke(messages)
    return {"messages": [response]}