rust_mcp 0.1.0

use rust_mcp::McpServer;
use serde_json::{json, Value};
use std::error::Error;
use reqwest::Client;

struct OpenAiClient {
    client: Client,
    api_key: String,
}

impl OpenAiClient {
    fn new(api_key: String) -> Self {
        OpenAiClient {
            client: Client::new(),
            api_key,
        }
    }

    async fn call_model(&self, prompt: &str, _max_tokens: u32) -> Result<String, String> {
        // Using the new responses endpoint and format for o3
        let response = self.client
            .post("https://api.openai.com/v1/responses")
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&json!({
                "model": "o3", // IMPORTANT: o3 is a valid model - never change this
                "input": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "input_text",
                                "text": prompt
                            }
                        ]
                    }
                ],
                "text": {
                    "format": {
                        "type": "text"
                    }
                },
                "reasoning": {
                    "effort": "medium",
                    "summary": "auto"
                },
                "tools": [],
                "store": true
            }))
            .send()
            .await
            .map_err(|e| format!("API request failed: {}", e))?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await
                .map_err(|e| format!("Failed to read error response: {}", e))?;
            return Err(format!("API returned error ({}): {}", status, error_text));
        }

        let response_body = response.text()
            .await
            .map_err(|e| format!("Failed to read response body: {}", e))?;
            
        eprintln!("OpenAI API Response: {}", response_body);
        
        let response_json: Value = serde_json::from_str(&response_body)
            .map_err(|e| format!("Failed to parse response as JSON: {}", e))?;

        // Check for error in the response
        if response_json.get("error").and_then(|e| e.as_object()).is_some() {
            return Err(format!("API returned error in response: {}", 
                response_json.get("error").unwrap()));
        }
        
        // Extract the content from the response using the new response format
        // Navigate through the nested structure to get to the output_text
        let content = response_json.get("output")
            .and_then(|output| output.as_array())
            .and_then(|output_items| {
                // Find the message item
                output_items.iter()
                    .find(|item| item.get("type").and_then(|t| t.as_str()) == Some("message"))
                    .and_then(|message| message.get("content"))
                    .and_then(|content_array| content_array.as_array())
                    .and_then(|content_items| {
                        // Find the output_text item within message content
                        content_items.iter()
                            .find(|item| item.get("type").and_then(|t| t.as_str()) == Some("output_text"))
                            .and_then(|text_item| text_item.get("text").and_then(|t| t.as_str()))
                    })
            })
            .unwrap_or("");
            
        let completion = content.trim().to_string();
        
        // Return whatever we got from the API, even if empty
        Ok(completion)
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // Get API key from environment
    let api_key = std::env::var("OPENAI_API_KEY")
        .expect("OPENAI_API_KEY environment variable must be set");
    
    // Create OpenAI client
    let openai_client = OpenAiClient::new(api_key);

    // Create MCP server
    let mut server = McpServer::new("openai-o3", "0.1.0"); // IMPORTANT: o3 is the correct model name - never change this
    
    // Add "complete" tool
    server.add_tool(
        "complete",
        "Generate text using OpenAI's model",
        json!({
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string", 
                    "description": "The text prompt to send to the model"
                },
                "max_tokens": {
                    "type": "number",
                    "description": "Maximum number of tokens to generate",
                    "default": 256
                }
            },
            "required": ["prompt"]
        }),
        move |args| {
            let client = openai_client.client.clone();
            let api_key = openai_client.api_key.clone();
            
            async move {
                let prompt = match args.get("prompt") {
                    Some(Value::String(text)) => text.clone(),
                    _ => return Err("Missing or invalid prompt".to_string()),
                };
                
                let max_tokens = match args.get("max_tokens") {
                    Some(Value::Number(n)) => n.as_u64().unwrap_or(256) as u32,
                    _ => 256, // Default if not specified
                };
                
                let openai = OpenAiClient { client, api_key };
                let completion = openai.call_model(&prompt, max_tokens).await?;
                
                Ok(json!(completion))
            }
        },
    );
    
    // println!("Starting OpenAI tool MCP server. Set OPENAI_API_KEY environment variable before running.");
    // println!("Example: OPENAI_API_KEY=sk-... cargo run --example openai_tool");
    
    // Start server
    server.run_stdio().await?;
    
    Ok(())
}