Make the LLM model and provider configurable

Author: sunu

.env.example

@@ -1,4 +1,12 @@
+# For OPENAI
 OPENAI_API_KEY="your-api-key"
+SMALL_MODEL_NAME="openai:gpt-4.1-mini"
+
+# For Mistral
+# MISTRAL_API_KEY="your-mistral-key"
+# SMALL_MODEL_NAME="mistral:mistral-medium-latest"
+
 GEODINI_API="https://geodini.k8s.labs.ds.io"
+
 STAC_CATALOG_NAME="planetarycomputer"
-STAC_CATALOG_URL="https://planetarycomputer.microsoft.com/api/stac/v1"
+STAC_CATALOG_URL="https://planetarycomputer.microsoft.com/api/stac/v1"

README.md

@@ -15,13 +15,20 @@ STAC Semantic Search allows users to find satellite imagery and geospatial data
 - **Visual Results**: Interactive maps showing search results and areas of interest
 
 
+## 🚧 Early Prototype
+
+This project is still in its very early stages and under active development. You might encounter some quirks or rough edges, but please give it a try and let us know what you think!
+
+
 ## 🚀 Getting Started
 
 ### Prerequisites
 
-- Python 3.12+
-- OpenAI API key
-- Docker (optional)
+- OpenAI API key / Anthropic API key / Mistral API key (or any other LLM [supported by Pydantic AI](https://ai.pydantic.dev/models/))
+- Docker and docker compose
+- A STAC catalog to search through
+- A [Geodini instance](https://github.com/sunu/geodini) to resolve geolocation queries
+
 
 ### Installation
 
@@ -55,7 +62,7 @@ STAC Semantic Search allows users to find satellite imagery and geospatial data
 ```bash
 curl -X POST "http://localhost:8000/search" \
      -H "Content-Type: application/json" \
-     -d '{"query": "Sentinel-2 imagery over France"}'
+     -d '{"query": "Sentinel-2 imagery"}'
 ```
 
 **Search Items**

helm-chart/Chart.yaml

@@ -1,18 +1,9 @@
 apiVersion: v2
-name: stac-search
-description: A Helm chart for STAC Semantic Search application
+name: stac-semantic-search
+description: A Helm chart for STAC Semantic Search API and Streamlit UI
 type: application
 version: 0.1.0
 appVersion: "0.1.0"
-keywords:
-  - stac
-  - search
-  - semantic
-  - geospatial
-  - satellite
-home: https://github.com/your-org/stac-semantic-search
+home: https://github.com/developmentseed/stac-semantic-search
 sources:
-  - https://github.com/your-org/stac-semantic-search
-maintainers:
-  - name: Your Name
-    email: [email protected]
+  - https://github.com/developmentseed/stac-semantic-search

stac_search/agents/collections_search.py

@@ -19,14 +19,12 @@
 # Constants
 MODEL_NAME = "all-MiniLM-L6-v2"
 DATA_PATH = os.environ.get("DATA_PATH", "data/chromadb")
-OPENAI_MODEL = "gpt-4o-mini"
 
 STAC_CATALOG_NAME = os.getenv("STAC_CATALOG_NAME", "planetarycomputer")
 STAC_COLLECTIONS_URL = os.getenv(
     "STAC_COLLECTIONS_URL", "https://planetarycomputer.microsoft.com/api/stac/v1"
 )
-
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", None)
+SMALL_MODEL_NAME = os.getenv("SMALL_MODEL_NAME", "openai:gpt-4.1-mini")
 
 
 @dataclass
@@ -49,7 +47,7 @@ class RankedCollections:
 
 
 rerank_agent = Agent(
-    "openai:gpt-4o-mini",
+    SMALL_MODEL_NAME,
     result_type=RankedCollections,
     deps_type=RerankContext,
     system_prompt="""

stac_search/agents/items_search.py

@@ -3,11 +3,12 @@
 import os
 from dataclasses import dataclass
 from pprint import pformat
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Union
 
 import requests
 from pydantic_ai import Agent, RunContext
 from pystac_client import Client
+from pydantic import BaseModel
 
 from stac_search.agents.collections_search import (
     collection_search,
@@ -16,6 +17,10 @@
 
 
 GEODINI_API = os.getenv("GEODINI_API", "https://geodini.k8s.labs.ds.io")
+SMALL_MODEL_NAME = os.getenv("SMALL_MODEL_NAME", "openai:gpt-4.1-mini")
+STAC_CATALOG_URL = os.getenv(
+    "STAC_CATALOG_URL", "https://planetarycomputer.microsoft.com/api/stac/v1"
+)
 
 logger = logging.getLogger(__name__)
 
@@ -38,7 +43,7 @@ class ItemSearchParams:
 
 
 search_items_agent = Agent(
-    "openai:gpt-4o-mini",
+    SMALL_MODEL_NAME,
     result_type=ItemSearchParams,
     deps_type=Context,
     system_prompt=f"""
@@ -60,7 +65,7 @@ class CollectionQuery:
 
 
 collection_query_framing_agent = Agent(
-    "openai:gpt-4o-mini",
+    SMALL_MODEL_NAME,
     result_type=CollectionQuery,
     system_prompt="""
 The user query is searching for relevant satellite imagery. 
@@ -108,7 +113,7 @@ class GeocodingResult:
 
 
 geocoding_agent = Agent(
-    "openai:gpt-4o-mini",
+    SMALL_MODEL_NAME,
     result_type=GeocodingResult,
     system_prompt="""
 For the given query, if it contains a location, return location query to be used to search for the location.
@@ -134,7 +139,7 @@ class TemporalRangeResult:
 
 
 temporal_range_agent = Agent(
-    "openai:gpt-4o-mini",
+    SMALL_MODEL_NAME,
     result_type=TemporalRangeResult,
     system_prompt="""
 For the given query, if it contains a temporal range, return the start date and end date. If it doesn't contain a temporal range, return None.
@@ -151,17 +156,42 @@ async def set_temporal_range(ctx: RunContext[Context]) -> TemporalRangeResult:
     return result.data
 
 
-@dataclass
-class Cql2Filter:
-    """Parameters to be used to query the STAC API"""
+class PropertyRef(BaseModel):
+    property: str
+
+
+class Geometry(BaseModel):
+    type: str
+    coordinates: Any
+
+
+class GeometryLiteral(BaseModel):
+    geometry: Geometry
+
 
+class PeriodLiteral(BaseModel):
+    period: List[str]
+
+
+FilterArg = Union[
+    "FilterExpr", PropertyRef, GeometryLiteral, PeriodLiteral, int, float, str
+]
+
+
+class FilterExpr(BaseModel):
     op: str
-    args: List[Any]
+    args: List[FilterArg]
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+FilterExpr.update_forward_refs()
 
 
 cql2_filter_agent = Agent(
-    "openai:gpt-4o-mini",
-    result_type=Cql2Filter,
+    SMALL_MODEL_NAME,
+    result_type=FilterExpr,
     system_prompt="""
 For the given query, construct a CQL2 filter to be used to query the STAC API only if required.
 Return None if the query doesn't require a CQL2 filter or if you can't determine the filter or if the property is not supported.
@@ -200,7 +230,7 @@ class Cql2Filter:
 
 
 @search_items_agent.tool
-async def construct_cql2_filter(ctx: RunContext[Context]) -> Cql2Filter:
+async def construct_cql2_filter(ctx: RunContext[Context]) -> FilterExpr | None:
     return await cql2_filter_agent.run(ctx.deps.query)
 
 
@@ -259,8 +289,7 @@ async def item_search(ctx: Context) -> ItemSearchResult:
         collections_to_search = default_target_collections
 
     # Actually perform the search
-    CATALOG_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
-    client = Client.open(CATALOG_URL)
+    client = Client.open(STAC_CATALOG_URL)
     params = {
         "max_items": 20,
         # looks like collections is required by Planetary Computer STAC API