Make things configurable by env vars

Author: sunu

stac_search/agents/collections_search.py

@@ -14,14 +14,13 @@
 
 # Constants
 MODEL_NAME = "all-MiniLM-L6-v2"
-DATA_PATH = "data/chromadb"
+DATA_PATH = os.environ.get("DATA_PATH", "data/chromadb")
 OPENAI_MODEL = "gpt-4o-mini"
 
-# STAC_CATALOG_NAME = "eoapi.dev"
-# STAC_COLLECTIONS_URL = "https://stac.eoapi.dev/"
-
-STAC_CATALOG_NAME = "planetarycomputer"
-STAC_COLLECTIONS_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
+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)
 

stac_search/agents/items_search.py

@@ -5,12 +5,15 @@
 import requests
 from pydantic_ai import Agent, RunContext
 from pprint import pprint
+import os
 
 from stac_search.agents.collections_search import (
     collection_search,
     CollectionWithExplanation,
 )
 
+GEODINI_API = os.getenv("GEODINI_API")
+
 
 @dataclass
 class Context:
@@ -197,7 +200,7 @@ async def construct_cql2_filter(ctx: RunContext[Context]) -> Cql2Filter:
 
 
 def get_polygon_from_geodini(location: str):
-    geodini_api = "https://api.geodini.labs.sunu.in/search"
+    geodini_api = f"{GEODINI_API}/search"
     response = requests.get(
         geodini_api,
         params={
@@ -275,10 +278,14 @@ async def item_search(ctx: Context) -> ItemSearchResult:
 
     if ctx.return_search_params_only:
         print("Returning STAC query parameters only")
-        return ItemSearchResult(search_params=params, aoi=polygon, explanation=explanation)
+        return ItemSearchResult(
+            search_params=params, aoi=polygon, explanation=explanation
+        )
 
     items = list(client.search(**params).items_as_dicts())
-    return ItemSearchResult(items=items, aoi=polygon, explanation=explanation, search_params=params)
+    return ItemSearchResult(
+        items=items, aoi=polygon, explanation=explanation, search_params=params
+    )
 
 
 async def main():

stac_search/load.py

@@ -5,15 +5,21 @@
 from sentence_transformers import SentenceTransformer
 import chromadb
 from pystac_client import Client
+import os
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Constants
 MODEL_NAME = "all-MiniLM-L6-v2"
-DATA_PATH = "data/chromadb"
+DATA_PATH = os.environ.get("DATA_PATH", "data/chromadb")
 
 
 def load_data(catalog_url, catalog_name):
     """Load STAC collections into the vector database"""
-    print("Initializing vector database...")
+    logger.info("Initializing vector database...")
 
     # Initialize the model
     model = SentenceTransformer(MODEL_NAME)
@@ -27,14 +33,14 @@ def load_data(catalog_url, catalog_name):
     # Initialize STAC client
     stac_client = Client.open(catalog_url)
 
-    print("Fetching STAC collections...")
+    logger.info("Fetching STAC collections...")
     collections = fetch_collections(stac_client)
-    print(f"Found {len(collections)} collections")
+    logger.info(f"Found {len(collections)} collections")
 
-    print("Generating embeddings and storing in vector database...")
+    logger.info("Generating embeddings and storing in vector database...")
     store_in_vector_db(collections, model, chroma_collection)
 
-    print("Data loading complete!")
+    logger.info("Data loading complete!")
 
 
 def fetch_collections(stac_client):
@@ -73,8 +79,15 @@ def store_in_vector_db(collections, model, chroma_collection):
 
 
 if __name__ == "__main__":
-    load_data(catalog_url="https://stac.eoapi.dev/", catalog_name="eoapi.dev")
-    load_data(
-        catalog_url="https://planetarycomputer.microsoft.com/api/stac/v1",
-        catalog_name="planetarycomputer",
+    # load_data(catalog_url="https://stac.eoapi.dev/", catalog_name="eoapi.dev")
+    # load_data(
+    #     catalog_url="https://planetarycomputer.microsoft.com/api/stac/v1",
+    #     catalog_name="planetarycomputer",
+    # )
+    import os
+
+    STAC_CATALOG_URL = os.environ.get(
+        "STAC_CATALOG_URL", "https://planetarycomputer.microsoft.com/api/stac/v1"
     )
+    STAC_CATALOG_NAME = os.environ.get("STAC_CATALOG_NAME", "planetarycomputer")
+    load_data(catalog_url=STAC_CATALOG_URL, catalog_name=STAC_CATALOG_NAME)