configure logging

Author: sunu

stac_search/__init__.py

@@ -1,3 +1,14 @@
 """
 STAC Natural Query - Vector search for STAC collections
 """
+
+import logging
+import os
+
+
+log_level = os.getenv("LOG_LEVEL", "INFO")
+
+logging.basicConfig(
+    level=log_level,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)

stac_search/agents/collections_search.py

@@ -7,11 +7,15 @@
 from sentence_transformers import SentenceTransformer
 import chromadb
 from typing import List, Dict, Any
-from pprint import pprint
+from pprint import pformat
 import time
+import logging
 
 from pydantic_ai import Agent
 
+
+logger = logging.getLogger(__name__)
+
 # Constants
 MODEL_NAME = "all-MiniLM-L6-v2"
 DATA_PATH = os.environ.get("DATA_PATH", "data/chromadb")
@@ -84,10 +88,10 @@ async def collection_search(
     # Initialize model and database connections
     model = SentenceTransformer(model_name)
     load_model_time = time.time()
-    print(f"Model loading time: {load_model_time - start_time:.4f} seconds")
+    logger.info(f"Model loading time: {load_model_time - start_time:.4f} seconds")
 
     client = chromadb.PersistentClient(path=data_path)
-    print(client.list_collections())
+    logger.info(client.list_collections())
     collection_name = f"{stac_catalog_name}_collections"
     collection = client.get_collection(name=collection_name)
 
@@ -122,7 +126,7 @@ async def collection_search(
 
 async def main():
     collections = await collection_search("Sentinel-2 imagery over France")
-    pprint(collections)
+    logger.info(pformat(collections))
 
 
 if __name__ == "__main__":

stac_search/agents/items_search.py

@@ -4,15 +4,19 @@
 from pystac_client import Client
 import requests
 from pydantic_ai import Agent, RunContext
-from pprint import pprint
+from pprint import pformat
 import os
+import logging
 
 from stac_search.agents.collections_search import (
     collection_search,
     CollectionWithExplanation,
 )
 
-GEODINI_API = os.getenv("GEODINI_API")
+
+GEODINI_API = os.getenv("GEODINI_API", "https://geodini.k8s.labs.ds.io")
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -87,9 +91,9 @@ class CollectionSearchResult:
 
 
 async def search_collections(query: str) -> CollectionSearchResult | None:
-    # print("Searching for relevant collections ...")
+    logger.info("Searching for relevant collections ...")
     collection_query = await collection_query_framing_agent.run(query)
-    print("Framed collection query: ", collection_query.data.query)
+    logger.info(f"Framed collection query: {collection_query.data.query}")
     if collection_query.data.is_specific:
         collections = await collection_search(collection_query.data.query)
         return CollectionSearchResult(collections=collections)
@@ -231,11 +235,11 @@ async def item_search(ctx: Context) -> ItemSearchResult:
     results = await search_items_agent.run(
         f"Find items for the query: {ctx.query}", deps=ctx
     )
-    print(results.data)
+    logger.info(results.data)
 
     # determine the collections to search
     target_collections = await search_collections(ctx.query) or []
-    print("Target collections: ", target_collections)
+    logger.info(f"Target collections: {pformat(target_collections)}")
     default_target_collections = [
         "landsat-8-c2-l2",
         "sentinel-2-l2a",
@@ -267,17 +271,15 @@ async def item_search(ctx: Context) -> ItemSearchResult:
         "filter": results.data.filter,
     }
 
-    print(f"Searching with params: {params}")
+    logger.info(f"Searching with params: {params}")
 
     polygon = get_polygon_from_geodini(results.data.location)
     if polygon:
-        print(f"Found polygon for {results.data.location}")
+        logger.info(f"Found polygon for {results.data.location}")
         params["intersects"] = polygon
 
-    # print(f"Searching with params: {params}")
-
     if ctx.return_search_params_only:
-        print("Returning STAC query parameters only")
+        logger.info("Returning STAC query parameters only")
         return ItemSearchResult(
             search_params=params, aoi=polygon, explanation=explanation
         )
@@ -291,7 +293,7 @@ async def item_search(ctx: Context) -> ItemSearchResult:
 async def main():
     ctx = Context(query="NAIP imagery from Washington state")
     results = await item_search(ctx)
-    pprint(results)
+    logger.info(pformat(results))
 
 
 if __name__ == "__main__":

stac_search/load.py

@@ -8,8 +8,7 @@
 import os
 import logging
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
+
 logger = logging.getLogger(__name__)
 
 # Constants