lakehouse/crates/aibridge/src/providers/ollama.rs

use async_trait::async_trait;
use crate::client::{AiClient, GenerateRequest, GenerateResponse, EmbedRequest, EmbedResponse};
use crate::provider::ProviderAdapter;

pub struct OllamaAdapter {
    client: AiClient,
}

impl OllamaAdapter {
    pub fn new(client: AiClient) -> Self {
        Self { client }
    }
}

#[async_trait]
impl ProviderAdapter for OllamaAdapter {
    fn name(&self) -> &str {
        "ollama"
    }

    async fn chat(&self, req: GenerateRequest) -> Result<GenerateResponse, String> {
        self.client.generate(req).await
    }

    async fn embed(&self, req: EmbedRequest) -> Result<EmbedResponse, String> {
        self.client.embed(req).await
    }

    async fn unload(&self, model: &str) -> Result<(), String> {
        let _: serde_json::Value = self.client.unload_model(model).await?;
        Ok(())
    }

    async fn health(&self) -> Result<serde_json::Value, String> {
        self.client.health().await
    }
}