MAF Ollama - llama3.2:3b

OllamaAgent.cs

using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Text.Json;
using Microsoft.Agents.AI;
using Microsoft.Extensions.AI;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol;
using OllamaSharp;
using OllamaSharp.Models.Chat;
using ChatMessage = Microsoft.Extensions.AI.ChatMessage;

namespace AgentService.Providers;

/// <summary>
/// AIAgent implementation using Ollama via OllamaSharp.
/// Handles tool calling by parsing JSON from model output (similar to FoundryLocalAgent).
/// </summary>
public class OllamaAgent : AIAgent
{
    private readonly string _ollamaEndpoint;
    private readonly string _model;
    private readonly string? _instructions;
    private readonly string? _name;
    private readonly string? _description;
    private readonly string? _id;
    private readonly IList<McpClientTool> _mcpTools;
    private readonly string _mcpToolsUrl;
    private readonly ILogger? _logger;
    private readonly OllamaApiClient _ollamaClient;
    private IHttpClientFactory _httpClientFactory;

    public OllamaAgent(
        string ollamaEndpoint,
        string model,
        string mcpToolsUrl,
        IList<McpClientTool> mcpTools,
        HttpClient httpClient,
        IHttpClientFactory httpClientFactory,
        string? instructions = null,
        string? name = null,
        string? description = null,
        ILogger? logger = null)
    {
        _ollamaEndpoint = ollamaEndpoint;
        _model = model;
        _mcpToolsUrl = mcpToolsUrl;
        _mcpTools = mcpTools;
        _instructions = instructions;
        _logger = logger;

        _name = name ?? "OllamaAgent";
        _description = description ?? "AI Agent powered by Ollama with MCP tools";
        _id = Guid.NewGuid().ToString();

        // Initialize OllamaSharp client
        _ollamaClient = new OllamaApiClient(new Uri(_ollamaEndpoint), _model);
        _httpClientFactory = httpClientFactory;
    }

    protected override string? IdCore => _id;
    public override string? Name => _name;
    public override string? Description => _description;

    public override AgentThread GetNewThread() => new OllamaAgentThread();

    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
    {
        return new OllamaAgentThread(serializedThread, jsonSerializerOptions);
    }

    public override async Task<AgentRunResponse> RunAsync(
        IEnumerable<ChatMessage> messages,
        AgentThread? thread = null,
        AgentRunOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        var agentThread = thread as OllamaAgentThread ?? new OllamaAgentThread();

        foreach (var msg in messages)
        {
            agentThread.MessageStore.Add(msg);
        }

        var allMessages = agentThread.MessageStore.ToList();

        var endpointUri = new Uri(_ollamaEndpoint);
        using var chatActivity = GenAITracing.StartChatSpan(
            model: _model,
            provider: GenAITracing.Providers.Ollama,
            serverAddress: endpointUri.Host,
            serverPort: endpointUri.Port);

        GenAITracing.SetMessages(chatActivity,
            inputMessages: allMessages.Select(m => new { role = m.Role.Value, content = m.Text }));
        GenAITracing.SetToolDefinitions(chatActivity,
            _mcpTools.Select(t => new { type = "function", name = t.Name, description = t.Description }));

        try
        {
            var result = await ProcessWithToolsAsync(allMessages, chatActivity, cancellationToken);

            var responseMessage = new ChatMessage(Microsoft.Extensions.AI.ChatRole.Assistant, result);
            agentThread.MessageStore.Add(responseMessage);

            GenAITracing.SetMessages(chatActivity,
                outputMessages: new[] { new { role = "assistant", content = result } });
            GenAITracing.SetResponseAttributes(chatActivity, responseModel: _model, finishReason: "stop");

            return new AgentRunResponse(responseMessage);
        }
        catch (Exception ex)
        {
            GenAITracing.RecordError(chatActivity, ex);
            throw;
        }
    }

    public override async IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(
        IEnumerable<ChatMessage> messages,
        AgentThread? thread = null,
        AgentRunOptions? options = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        var response = await RunAsync(messages, thread, options, cancellationToken);

        foreach (var msg in response.Messages)
        {
            yield return new AgentRunResponseUpdate(msg.Role, msg.Text);
        }
    }

    private async Task<string> ProcessWithToolsAsync(
        IList<ChatMessage> chatMessages,
        Activity? parentActivity,
        CancellationToken cancellationToken)
    {
        var toolDefs = BuildToolDefs();
        var jsonExample = @"{""name"": ""tool_name"", ""arguments"": {""param"": ""value""}}";

        var systemPrompt = _mcpTools.Count > 0
            ? $"""
                {_instructions ?? "You are a helpful assistant with access to tools."}
                
                When you need to use a tool, respond ONLY with a JSON object in this exact format:
                {jsonExample}

                Available tools:
                {toolDefs}
                
                If you don't need a tool, respond normally with text.
                """
            : _instructions ?? "You are a helpful assistant.";

        // Build Ollama chat messages
        var ollamaMessages = new List<OllamaSharp.Models.Chat.Message>
        {
            new() { Role = OllamaSharp.Models.Chat.ChatRole.System, Content = systemPrompt }
        };

        foreach (var msg in chatMessages)
        {
            ollamaMessages.Add(new OllamaSharp.Models.Chat.Message
            {
                Role = msg.Role == Microsoft.Extensions.AI.ChatRole.User
                    ? OllamaSharp.Models.Chat.ChatRole.User
                    : OllamaSharp.Models.Chat.ChatRole.Assistant,
                Content = msg.Text ?? ""
            });
        }

        _logger?.LogDebug("[OllamaAgent] Sending request to model {Model}...", _model);

        var chatRequest = new OllamaSharp.Models.Chat.ChatRequest
        {
            Model = _model,
            Messages = ollamaMessages,
            Stream = false,
            Tools = _mcpTools.Select(t => new
            {
                name = t.Name,
                description = t.Description,
                parameters = t.JsonSchema.GetRawText()
            }).ToList()
        };

        // Consume the async enumerable to get the final response
        var content = "";
        await foreach (var chunk in _ollamaClient.ChatAsync(chatRequest, cancellationToken))
        {
            if (chunk?.Message?.Content is not null)
            {
                content += chunk.Message.Content;
            }
        }

        _logger?.LogDebug("[OllamaAgent] Model response: {Content}", content);

        // Parse tool calls from content
        var toolCalls = ToolCallParser.Parse(content);
        var validToolCalls = toolCalls
            .Where(tc => _mcpTools.Any(t => t.Name == tc.Name))
            .ToList();

        if (validToolCalls.Count == 0)
        {
            return content;
        }

        _logger?.LogInformation("[OllamaAgent] Executing {Count} tool(s): {Tools}",
            validToolCalls.Count,
            string.Join(", ", validToolCalls.Select(tc => tc.Name)));

        var toolResults = await ExecuteToolsAsync(validToolCalls, cancellationToken);

        var toolResultsText = string.Join("\n\n", toolResults.Select(r =>
            r.Error is null
                ? $"Tool '{r.Name}' result:\n{r.Result}"
                : $"Tool '{r.Name}' error: {r.Error}"));

        // Add tool results to conversation and get final response
        ollamaMessages.Add(new OllamaSharp.Models.Chat.Message
        {
            Role = OllamaSharp.Models.Chat.ChatRole.Assistant,
            Content = toolResultsText
        });
        ollamaMessages.Add(new OllamaSharp.Models.Chat.Message
        {
            Role = OllamaSharp.Models.Chat.ChatRole.User,
            Content = $"Tool execution results:\n{toolResultsText}\n\nPlease provide your final response based on these results."
        });

        var finalRequest = new OllamaSharp.Models.Chat.ChatRequest
        {
            Model = _model,
            Messages = ollamaMessages,
            Stream = false
        };

        // Consume the async enumerable to get the final response
        var finalContent = "";
        await foreach (var chunk in _ollamaClient.ChatAsync(finalRequest, cancellationToken))
        {
            if (chunk?.Message?.Content is not null)
            {
                finalContent += chunk.Message.Content;
            }
        }
        
        return string.IsNullOrEmpty(finalContent) ? toolResultsText : finalContent;
    }

    private async Task<List<ToolExecutionResult>> ExecuteToolsAsync(
        List<ToolCall> toolCalls,
        CancellationToken cancellationToken)
    {
        var toolTasks = toolCalls.Select(async toolCall =>
        {
            using var toolActivity = GenAITracing.StartToolSpan(
                toolName: toolCall.Name,
                toolCallId: Guid.NewGuid().ToString("N")[..12],
                arguments: toolCall.Arguments);

            try
            {
                var result = await CallMcpToolAsync(toolCall.Name, toolCall.Arguments, cancellationToken);
                GenAITracing.SetToolResult(toolActivity, result);
                return new ToolExecutionResult(toolCall.Name, result ?? "", null);
            }
            catch (Exception ex)
            {
                GenAITracing.RecordError(toolActivity, ex);
                return new ToolExecutionResult(toolCall.Name, "", ex.Message);
            }
        });

        return (await Task.WhenAll(toolTasks)).ToList();
    }

    private record ToolExecutionResult(string Name, string Result, string? Error);

    private string BuildToolDefs()
    {
        return string.Join("\n\n", _mcpTools.Select(tool =>
        {
            var paramsDesc = tool.JsonSchema.TryGetProperty("properties", out var props)
                ? string.Join(", ", props.EnumerateObject().Select(p => $"{p.Name}: {GetTypeDescription(p.Value)}"))
                : "no parameters";

            return $"- {tool.Name}: {tool.Description}\n  Parameters: {paramsDesc}";
        }));
    }

    private static string GetTypeDescription(JsonElement element)
    {
        if (element.TryGetProperty("type", out var typeEl))
            return typeEl.GetString() ?? "any";
        return "any";
    }

    private async Task<string> CallMcpToolAsync(
        string toolName,
        Dictionary<string, JsonElement> arguments,
        CancellationToken cancellationToken)
    {
        var transport = new HttpClientTransport(new HttpClientTransportOptions
        {
            Endpoint = new Uri(_mcpToolsUrl)
        }, _httpClientFactory.CreateClient(), ownsHttpClient: false);

        await using var mcpClient = await McpClient.CreateAsync(transport, cancellationToken: cancellationToken);

        var convertedArgs = arguments.ToDictionary(
            kvp => kvp.Key,
            kvp => ConvertJsonElement(kvp.Value));

        var result = await mcpClient.CallToolAsync(toolName, convertedArgs, cancellationToken: cancellationToken);

        var textContent = result.Content.OfType<TextContentBlock>().FirstOrDefault();
        return textContent?.Text ?? "Tool returned no content";
    }

    private static object? ConvertJsonElement(JsonElement element) => element.ValueKind switch
    {
        JsonValueKind.String => element.GetString(),
        JsonValueKind.Number => element.TryGetInt64(out var l) ? l : element.GetDouble(),
        JsonValueKind.True => true,
        JsonValueKind.False => false,
        JsonValueKind.Null => null,
        _ => element.GetRawText()
    };
}