feat: support image input to LLM clients

Cargo.toml

@@ -115,3 +115,5 @@ aws-config = "1.6.2"
 aws-sdk-s3 = "1.85.0"
 aws-sdk-sqs = "1.67.0"
 numpy = "0.25.0"
+infer = "0.19.0"
+serde_with = { version = "3.13.0", features = ["base64"] }

examples/image_search/README.md

@@ -13,6 +13,7 @@ We appreciate a star ⭐ at [CocoIndex Github](https://github.com/cocoindex-io/c
 - CLIP ViT-L/14 - Embeddings Model for images and query
 - Qdrant for Vector Storage
 - FastApi for backend
+- Ollama (Optional) for generating image captions using `gemma3`.
 
 ## Setup
 - Make sure Postgres and Qdrant are running
@@ -21,7 +22,15 @@ We appreciate a star ⭐ at [CocoIndex Github](https://github.com/cocoindex-io/c
   export COCOINDEX_DATABASE_URL="postgres://cocoindex:cocoindex@localhost/cocoindex"
   ```
 
-## Run
+## (Optional) Run Ollama
+
+- This enables automatic image captioning
+```
+ollama pull gemma3
+ollama serve
+```
+
+## Run the App
 - Install dependencies:
   ```
   pip install -e .

examples/image_search/main.py

@@ -6,6 +6,7 @@
 from typing import Any, Literal
 
 import cocoindex
+import requests
 import torch
 from dotenv import load_dotenv
 from fastapi import FastAPI, Query
@@ -17,10 +18,21 @@
 
 QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6334/")
 QDRANT_COLLECTION = "ImageSearch"
+OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434/")
+OLLAMA_MODEL = "gemma3"
 CLIP_MODEL_NAME = "openai/clip-vit-large-patch14"
 CLIP_MODEL_DIMENSION = 768
 
 
+def ollama_has_model(model) -> bool:
+    try:
+        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=1)
+        r.raise_for_status()
+        return any(m.get("name") == model for m in r.json().get("models", []))
+    except Exception:
+        return False
+
+
 @functools.cache
 def get_clip_model() -> tuple[CLIPModel, CLIPProcessor]:
     model = CLIPModel.from_pretrained(CLIP_MODEL_NAME)
@@ -69,12 +81,49 @@ def image_object_embedding_flow(
     )
     img_embeddings = data_scope.add_collector()
     with data_scope["images"].row() as img:
+        has_gemma3 = ollama_has_model(OLLAMA_MODEL)
+        if has_gemma3:
+            img["caption"] = flow_builder.transform(
+                cocoindex.functions.ExtractByLlm(
+                    llm_spec=cocoindex.llm.LlmSpec(
+                        api_type=cocoindex.LlmApiType.OLLAMA, model=OLLAMA_MODEL
+                    ),
+                    # Replace by this spec below, to use OpenAI API model instead of ollama
+                    #   llm_spec=cocoindex.LlmSpec(
+                    #       api_type=cocoindex.LlmApiType.OPENAI, model="gpt-4o"),
+                    # Replace by this spec below, to use Gemini API model
+                    #   llm_spec=cocoindex.LlmSpec(
+                    #       api_type=cocoindex.LlmApiType.GEMINI, model="gemini-2.0-flash"),
+                    # Replace by this spec below, to use Anthropic API model
+                    #   llm_spec=cocoindex.LlmSpec(
+                    #       api_type=cocoindex.LlmApiType.ANTHROPIC, model="claude-3-5-sonnet-latest"),
+                    instruction=(
+                        "Describe the image in one detailed sentence. "
+                        "Name all visible animal species, objects, and the main scene. "
+                        "Be specific about type, color, and notable features. "
+                        "Mention what each animal is doing."
+                    ),
+                    output_type=str,
+                ),
+                image=img["content"],
+            )
         img["embedding"] = img["content"].transform(embed_image)
-        img_embeddings.collect(
-            id=cocoindex.GeneratedField.UUID,
-            filename=img["filename"],
-            embedding=img["embedding"],
-        )
+
+        collect_fields = {
+            "id": cocoindex.GeneratedField.UUID,
+            "filename": img["filename"],
+            "embedding": img["embedding"],
+        }
+
+        if has_gemma3:
+            print(
+                f"Ollama model '{OLLAMA_MODEL}' is available — captions will be extracted."
+            )
+            collect_fields["caption"] = img["caption"]
+        else:
+            print(f"Ollama model '{OLLAMA_MODEL}' not found — skipping captioning.")
+
+        img_embeddings.collect(**collect_fields)
 
     img_embeddings.export(
         "img_embeddings",
@@ -126,11 +175,18 @@ def search(
         collection_name=QDRANT_COLLECTION,
         query_vector=("embedding", query_embedding),
         limit=limit,
+        with_payload=True,
     )
 
     return {
         "results": [
-            {"filename": result.payload["filename"], "score": result.score}
+            {
+                "filename": result.payload["filename"],
+                "score": result.score,
+                "caption": result.payload.get(
+                    "caption"
+                ),  # Include caption if available
+            }
             for result in search_results
         ]
     }

examples/image_search/pyproject.toml

@@ -11,6 +11,7 @@ dependencies = [
     "transformers>=4.29.0",
     "qdrant-client>=1.14.2",
     "uvicorn>=0.34.3",
+    "requests>=2.32.4",
 ]
 
 [tool.setuptools]

src/llm/anthropic.rs

@@ -1,8 +1,10 @@
 use crate::llm::{
-    LlmGenerateRequest, LlmGenerateResponse, LlmGenerationClient, OutputFormat, ToJsonSchemaOptions,
+    LlmGenerateRequest, LlmGenerateResponse, LlmGenerationClient, OutputFormat,
+    ToJsonSchemaOptions, detect_image_mime_type,
 };
 use anyhow::{Context, Result, bail};
 use async_trait::async_trait;
+use base64::prelude::*;
 use json5;
 use serde_json::Value;
 
@@ -36,9 +38,31 @@ impl LlmGenerationClient for Client {
         &self,
         request: LlmGenerateRequest<'req>,
     ) -> Result<LlmGenerateResponse> {
+        let mut user_content_parts: Vec<serde_json::Value> = Vec::new();
+
+        // Add image part if present
+        if let Some(image_bytes) = &request.image {
+            let base64_image = BASE64_STANDARD.encode(image_bytes.as_ref());
+            let mime_type = detect_image_mime_type(image_bytes.as_ref())?;
+            user_content_parts.push(serde_json::json!({
+                "type": "image",
+                "source": {
+                    "type": "base64",
+                    "media_type": mime_type,
+                    "data": base64_image,
+                }
+            }));
+        }
+
+        // Add text part
+        user_content_parts.push(serde_json::json!({
+            "type": "text",
+            "text": request.user_prompt
+        }));
+
         let messages = vec![serde_json::json!({
             "role": "user",
-            "content": request.user_prompt
+            "content": user_content_parts
         })];
 
         let mut payload = serde_json::json!({

src/llm/gemini.rs

@@ -2,8 +2,9 @@ use crate::prelude::*;
 
 use crate::llm::{
     LlmEmbeddingClient, LlmGenerateRequest, LlmGenerateResponse, LlmGenerationClient, OutputFormat,
-    ToJsonSchemaOptions,
+    ToJsonSchemaOptions, detect_image_mime_type,
 };
+use base64::prelude::*;
 use phf::phf_map;
 use serde_json::Value;
 use urlencoding::encode;
@@ -70,10 +71,27 @@ impl LlmGenerationClient for Client {
         &self,
         request: LlmGenerateRequest<'req>,
     ) -> Result<LlmGenerateResponse> {
-        // Compose the prompt/messages
+        let mut user_parts: Vec<serde_json::Value> = Vec::new();
+
+        // Add text part first
+        user_parts.push(serde_json::json!({ "text": request.user_prompt }));
+
+        // Add image part if present
+        if let Some(image_bytes) = &request.image {
+            let base64_image = BASE64_STANDARD.encode(image_bytes.as_ref());
+            let mime_type = detect_image_mime_type(image_bytes.as_ref())?;
+            user_parts.push(serde_json::json!({
+                "inlineData": {
+                    "mimeType": mime_type,
+                    "data": base64_image
+                }
+            }));
+        }
+
+        // Compose the contents
         let contents = vec![serde_json::json!({
             "role": "user",
-            "parts": [{ "text": request.user_prompt }]
+            "parts": user_parts
         })];
 
         // Prepare payload

src/llm/mod.rs

@@ -1,9 +1,12 @@
 use crate::prelude::*;
 
 use crate::base::json_schema::ToJsonSchemaOptions;
+use infer::Infer;
 use schemars::schema::SchemaObject;
 use std::borrow::Cow;
 
+static INFER: LazyLock<Infer> = LazyLock::new(Infer::new);
+
 #[derive(Debug, Clone, Copy, Serialize, Deserialize)]
 pub enum LlmApiType {
     Ollama,
@@ -36,6 +39,7 @@ pub struct LlmGenerateRequest<'a> {
     pub model: &'a str,
     pub system_prompt: Option<Cow<'a, str>>,
     pub user_prompt: Cow<'a, str>,
+    pub image: Option<Cow<'a, [u8]>>,
     pub output_format: Option<OutputFormat<'a>>,
 }
 
@@ -141,3 +145,11 @@ pub fn new_llm_embedding_client(
     };
     Ok(client)
 }
+
+pub fn detect_image_mime_type(bytes: &[u8]) -> Result<&'static str> {
+    let infer = &*INFER;
+    match infer.get(bytes) {
+        Some(info) if info.mime_type().starts_with("image/") => Ok(info.mime_type()),
+        _ => bail!("Unknown or unsupported image format"),
+    }
+}

src/llm/ollama.rs

@@ -2,6 +2,7 @@ use crate::prelude::*;
 
 use super::LlmGenerationClient;
 use schemars::schema::SchemaObject;
+use serde_with::{base64::Base64, serde_as};
 
 pub struct Client {
     generate_url: String,
@@ -14,10 +15,13 @@ enum OllamaFormat<'a> {
     JsonSchema(&'a SchemaObject),
 }
 
+#[serde_as]
 #[derive(Debug, Serialize)]
 struct OllamaRequest<'a> {
     pub model: &'a str,
     pub prompt: &'a str,
+    #[serde_as(as = "Option<Vec<Base64>>")]
+    pub images: Option<Vec<&'a [u8]>>,
     pub format: Option<OllamaFormat<'a>>,
     pub system: Option<&'a str>,
     pub stream: Option<bool>,
@@ -52,6 +56,7 @@ impl LlmGenerationClient for Client {
         let req = OllamaRequest {
             model: request.model,
             prompt: request.user_prompt.as_ref(),
+            images: request.image.as_deref().map(|img| vec![img.as_ref()]),
             format: request.output_format.as_ref().map(
                 |super::OutputFormat::JsonSchema { schema, .. }| {
                     OllamaFormat::JsonSchema(schema.as_ref())

src/llm/openai.rs

@@ -1,18 +1,21 @@
 use crate::api_bail;
 
-use super::{LlmEmbeddingClient, LlmGenerationClient};
+use super::{LlmEmbeddingClient, LlmGenerationClient, detect_image_mime_type};
 use anyhow::Result;
 use async_openai::{
     Client as OpenAIClient,
     config::OpenAIConfig,
     types::{
-        ChatCompletionRequestMessage, ChatCompletionRequestSystemMessage,
+        ChatCompletionRequestMessage, ChatCompletionRequestMessageContentPartImage,
+        ChatCompletionRequestMessageContentPartText, ChatCompletionRequestSystemMessage,
         ChatCompletionRequestSystemMessageContent, ChatCompletionRequestUserMessage,
-        ChatCompletionRequestUserMessageContent, CreateChatCompletionRequest,
-        CreateEmbeddingRequest, EmbeddingInput, ResponseFormat, ResponseFormatJsonSchema,
+        ChatCompletionRequestUserMessageContent, ChatCompletionRequestUserMessageContentPart,
+        CreateChatCompletionRequest, CreateEmbeddingRequest, EmbeddingInput, ImageDetail,
+        ResponseFormat, ResponseFormatJsonSchema,
     },
 };
 use async_trait::async_trait;
+use base64::prelude::*;
 use phf::phf_map;
 
 static DEFAULT_EMBEDDING_DIMENSIONS: phf::Map<&str, u32> = phf_map! {
@@ -64,11 +67,32 @@ impl LlmGenerationClient for Client {
         }
 
         // Add user message
+        let user_message_content = match request.image {
+            Some(img_bytes) => {
+                let base64_image = BASE64_STANDARD.encode(img_bytes.as_ref());
+                let mime_type = detect_image_mime_type(img_bytes.as_ref())?;
+                let image_url = format!("data:{};base64,{}", mime_type, base64_image);
+                ChatCompletionRequestUserMessageContent::Array(vec![
+                    ChatCompletionRequestUserMessageContentPart::Text(
+                        ChatCompletionRequestMessageContentPartText {
+                            text: request.user_prompt.into_owned(),
+                        },
+                    ),
+                    ChatCompletionRequestUserMessageContentPart::ImageUrl(
+                        ChatCompletionRequestMessageContentPartImage {
+                            image_url: async_openai::types::ImageUrl {
+                                url: image_url,
+                                detail: Some(ImageDetail::Auto),
+                            },
+                        },
+                    ),
+                ])
+            }
+            None => ChatCompletionRequestUserMessageContent::Text(request.user_prompt.into_owned()),
+        };
         messages.push(ChatCompletionRequestMessage::User(
             ChatCompletionRequestUserMessage {
-                content: ChatCompletionRequestUserMessageContent::Text(
-                    request.user_prompt.into_owned(),
-                ),
+                content: user_message_content,
                 ..Default::default()
             },
         ));