client.py

from langgraph.graph import StateGraph, START
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI
from typing import TypedDict, Annotated
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
import asyncio
from langchain_mcp_adapters.client import MultiServerMCPClient
import os
from typing import Optional
import sys

api_key = os.getenv("GEMINI_API_KEY")

llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    temperature=0.2,
    google_api_key=api_key
)

print("🔍 DEBUG: MCP Server Path Resolution")
print("=" * 60)
print(f"📂 Current working directory: {os.getcwd()}")
print(f"📂 __file__ is: {__file__}")
print(f"📂 __file__ directory: {os.path.dirname(__file__)}")

# Try multiple path strategies
server_path_options = [
    os.path.join(os.path.dirname(__file__), "server.py"),
    os.path.join(os.getcwd(), "server.py"),
    os.path.abspath("server.py"),
    "server.py"
]

server_path = None
for path in server_path_options:
    print(f"🔍 Checking: {path} ... ", end="")
    if os.path.exists(path):
        print("✅ FOUND!")
        server_path = path
        break
    else:
        print("❌ Not found")

if not server_path:
    print("\n❌ ERROR: server.py not found in any expected location!")
    print(f"📁 Files in current directory: {os.listdir('.')}")
    if os.path.dirname(__file__):
        print(f"📁 Files in __file__ directory: {os.listdir(os.path.dirname(__file__))}")
    raise FileNotFoundError("server.py not found. Make sure it's uploaded to your HF Space.")
else:
    print(f"✅ Using server.py at: {server_path}")

print(f"🐍 Python executable: {sys.executable}")
print("=" * 60)

# Initialize MCP Client
try:
    # In client.py - pass environment explicitly
    client = MultiServerMCPClient(
    {
        "FistalMCP": {
            "transport": "stdio",
            "command": sys.executable,
            "args": ["-u", server_path],
            "env": {  # Explicitly pass environment variables
                **os.environ,  # Pass all current environment
                "GROQ_API_KEY": os.getenv("GROQ_API_KEY", ""),
                "HF_TOKEN": os.getenv("HF_TOKEN", ""),
                "GOOGLE_API_KEY_1": os.getenv("GOOGLE_API_KEY_1", ""),
                "GOOGLE_API_KEY_2": os.getenv("GOOGLE_API_KEY_2", ""),
                "GOOGLE_API_KEY_3": os.getenv("GOOGLE_API_KEY_3", ""),
                "PATH": os.getenv("PATH", ""),
            }
        }
    }
)
    print("✅ MCP Client initialized successfully!")
except Exception as e:
    print(f"❌ Failed to initialize MCP Client: {e}")
    import traceback
    traceback.print_exc()
    raise


class ChatState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    dataset_topic: str
    num_samples: int
    model_name: str
    task_type: str
    dataset_path: Optional[str]
    converted_path: Optional[str]
    model_path: Optional[str]
    hf_url: Optional[str]


async def my_graph():
    """Agent graph that handles mcp tools"""
    tools = await client.get_tools()

    available_tools = []
    tool_order = ["generate_json_data", "format_json", "finetune_model", "llm_as_judge"]
    available_tools = []
    for tool_name in tool_order:
        for tool in tools:
            if tool.name == tool_name:
                available_tools.append(tool)
                break

    llm_toolkit = llm.bind_tools(available_tools)

    async def chat_node(state: ChatState):
        messages = state["messages"]
        dataset_topic = state['dataset_topic']
        if isinstance(dataset_topic, list):
            dataset_topic = dataset_topic[0] if dataset_topic else "unknown"
        
        num_samples = state['num_samples']
        if isinstance(num_samples, list):
            num_samples = num_samples[0] if num_samples else 100
        
        model_name = state['model_name']
        if isinstance(model_name, list):
            model_name = model_name[0] if model_name else "unknown"
        
        task_type = state['task_type']
        if isinstance(task_type, list):
            task_type = task_type[0] if task_type else "text-generation"
        
        system_msg = f"""You are Fistal, an AI fine-tuning assistant.

**User's Configuration:**
- Dataset Topic: {dataset_topic}
- Number of Samples: {num_samples}
- Model to Fine-tune: {model_name}
- Task Type: {task_type}
- Evaluation : Using LLM

**Your Workflow:**
1. Use generate_json_data with topic="{dataset_topic}", task_type="{task_type}", num_samples={num_samples}
   - This returns a dictionary with a "data" field containing the raw dataset
   
2. Use format_json with the "data" field from step 1
   - Pass: raw_data=<the data list from step 1>
   - This returns a dictionary with a "data" field containing formatted data
   
3. Use finetune_model with the "data" field from step 2 and model_name="{model_name}"
   - Pass: formatted_data=<the data list from step 2>, model_name="{model_name}"
   - This returns the Hugging Face repo URL
   
4. Use llm_as_judge with the repo_id from step 3
   - Pass: repo_id=<the HF repo from step 3>, topic="{dataset_topic}", task_type="{task_type}"

**FINAL STEP - CRITICAL:**
5. After completing all tools, you MUST return:
   - The Hugging Face model URL from step 3
   - The evaluation report from step 4
   - Format your final response as:
   
🎉 **Fine-tuning Complete!**

**🤗 Model Repository:** [HF Repo Link] \n\n
**📊 Evaluation Report:** [Full report from llm_as_judge]

**IMPORTANT:**
- Tools pass DATA directly, not file paths
- Always mention the tool you are going to use first and then proceed with the tool action
- Extract the "data" field from each tool's response and pass it to the next tool
- After llm_as_judge completes, return both the HF URL and evaluation report
- Keep the user informed of progress at each step
- If a step takes time, do not stay idle. Inform users about short interesting facts about Modal, Unsloth, Gemini, Gradio , HuggingFace and MCP, do not repeat them.
- Try to add atleast 1 new fact every 10 seconds.
- Report any errors clearly
- Do not mention internal data structures or file paths"""

    
        full_messages = [SystemMessage(content=system_msg)] + messages
        response = await llm_toolkit.ainvoke(full_messages)
        return {'messages': [response]}
    
    tool_node = ToolNode(available_tools)

    graph = StateGraph(ChatState)

    graph.add_node("chat_node", chat_node)
    graph.add_node("tools", tool_node)

    graph.add_edge(START, "chat_node")
    graph.add_conditional_edges("chat_node", tools_condition)
    graph.add_edge("tools", "chat_node")

    chat = graph.compile()

    return chat




async def run_fistal(
    dataset_topic: str,
    num_samples: int,
    model_name: str,
    task_type: str
):
    chatbot = await my_graph()
    user_message = f"""Execute the complete fine-tuning workflow:
- Generate {num_samples} training examples about {dataset_topic}
- Fine-tune {model_name}
- Evaluate for {task_type} task

Start now!"""
    initial_state = {
        "messages": [HumanMessage(content=user_message)],
        "dataset_topic": dataset_topic,
        "num_samples": num_samples,
        "model_name": model_name,
        "task_type": task_type,
        "dataset_path": None,
        "converted_path": None,
        "model_path": None,
        "hf_url": None
    }
    facts = {
        "generate_json_data": [
            "💡 Using parallel batch generation with multiple API keys for 3x speed!",
            "📊 Quality over quantity - diverse examples lead to better models!",
            "🎯 Generating diverse prompt-response pairs...",
        ],
        "format_json": [
            "🔄 Converting to chat format optimized for instruction tuning...",
            "💬 Proper formatting helps models understand conversation structure!",
            "🎨 Applying ChatML format for consistency...",
            "✅ Validating JSON structure for training compatibility...",
            "🔧 Optimizing token distribution across examples..."
        ],
        "finetune_model": [
            "🏋️ Training on Modal's serverless T4 GPU...",
            "💡 Using 4-bit quantization to fit in 16GB VRAM!",
            "🦥 Unsloth makes training 2x faster with 70% less memory!",
            "⚡ LoRA fine-tuning updates only 0.1% of model parameters!",
            "🎯 Typical training time: 10-20 minutes for 500 samples...",
            "🔥 Your model is learning patterns from authentic data!",
            "☁️ Uploading to HuggingFace - your model will be public soon!"
        ],
        "llm_as_judge": [
            "📊 Generating evaluation test cases...",
            "🤖 LLM-as-judge provides qualitative insights!",
            "✨ Testing model coherence, relevance, and accuracy...",
            "📝 Creating comprehensive evaluation report...",
            "🔍 Analyzing response quality and task alignment...",
            "📝 Creating comprehensive evaluation report...",
            "📈 Comparing outputs against expected responses...",
            "🎯 Assessing model's understanding of the domain...",
            "✅ Finalizing evaluation metrics.."
        ]
    }

    current_tool = None
    fact_i = 0

    async for event in chatbot.astream(initial_state):
        if "tools" in event:
            messages = event["tools"].get("messages", [])
            for msg in messages:
                if hasattr(msg,"name"):
                    tool_name = msg.name
                    current_tool = tool_name
                    fact_i = 0
                    yield f"\n{'-'*60}\n"
                    yield f"🔄 **Using: {tool_name}**\n\n"
                    if tool_name in facts:
                        yield f"{facts[tool_name][0]}\n"
                        await asyncio.sleep(0.3)

        if "chat_node" in event:
            messages = event["chat_node"].get("messages", [])
            for msg in messages:
                if hasattr(msg, 'content') and msg.content:
                    raw_content = msg.content
                    content = ""
                    
                    if isinstance(raw_content, list):
                        for item in raw_content:
                            if isinstance(item, dict) and item.get('type') == 'text':
                                content += item.get('text', '')
                        content = content.strip()
                    elif isinstance(raw_content, str):
                        content = raw_content
                    else:
                        content = str(raw_content)
                    
                    if content and len(content) > 20 and "tool_calls" not in content.lower():
                        yield f"\n🤖 **Fistal:** {content}\n"
                    
                    if current_tool and current_tool in facts:
                        fact_i += 1
                        if fact_i < len(facts[current_tool]):
                            yield f"\n💡 {facts[current_tool][fact_i]}\n"
                            await asyncio.sleep(0.3)
    yield "✅ **Successfully finetuned!**\n"



async def main():
    """Test the agent. Only for running client.py"""
    print("Testing Fistal Agent\n")
    
    result = await run_fistal(
        "python programming",
        5,
        "unsloth/Llama-3.2-1B-Instruct-bnb-4bit",
        "text-generation"
    )
    
    print(f"\nAgent Response:\n{result}")


if __name__ == '__main__':
    asyncio.run(main())