Merge pull request #17 from developmentseed/log

Log all the steps into aimstack

Author: srmsoumya

.gitignore

@@ -160,4 +160,7 @@ cython_debug/
 #.idea/
 
 # streamlit cache
-cache/
+cache/
+
+# AIM experiment runs
+.aim/

agents/l4m_agent.py

@@ -18,7 +18,7 @@ def base_agent(
         llm=llm,
         tools=tools,
         agent=agent_type,
-        max_iterations=2,
+        max_iterations=5,
         early_stopping_method="generate",
         verbose=True,
         # memory=memory,

app.py

@@ -5,11 +5,13 @@
 import streamlit as st
 from streamlit_folium import folium_static
 
+import langchain
 from langchain.agents import AgentType
 from langchain.chat_models import ChatOpenAI
 from langchain.tools import Tool, DuckDuckGoSearchRun
 from langchain.callbacks import (
     StreamlitCallbackHandler,
+    AimCallbackHandler,
     get_openai_callback,
 )
 
@@ -21,6 +23,9 @@
 from tools.stac.search import STACSearchTool
 from agents.l4m_agent import base_agent
 
+# DEBUG
+langchain.debug = True
+
 
 @st.cache_resource(ttl="1h")
 def get_agent(
@@ -60,8 +65,6 @@ def get_agent(
 
 
 def run_query(agent, query):
-    st_callback = StreamlitCallbackHandler(st.container())
-    response = agent.run(query, callbacks=[st_callback])
     return response
 
 
@@ -117,7 +120,9 @@ def plot_vector(df):
     st.session_state.total_cost = 0
 
 with st.sidebar:
-    openai_api_key = st.text_input("OpenAI API Key", type="password")
+    openai_api_key = os.getenv("OPENAI_API_KEY")
+    if not openai_api_key:
+        openai_api_key = st.text_input("OpenAI API Key", type="password")
 
     st.subheader("OpenAI Usage")
     total_tokens = st.empty()
@@ -147,9 +152,18 @@ def plot_vector(df):
         st.info("Please add your OpenAI API key to continue.")
         st.stop()
 
+    aim_callback = AimCallbackHandler(
+        repo=".",
+        experiment_name="LLLLLM: Base Agent v0.1",
+    )
+
     agent = get_agent(openai_api_key)
+
     with get_openai_callback() as cb:
-        response = run_query(agent, prompt)
+        st_callback = StreamlitCallbackHandler(st.container())
+        response = agent.run(prompt, callbacks=[st_callback, aim_callback])
+
+        aim_callback.flush_tracker(langchain_asset=agent, reset=False, finish=True)
 
         # Log OpenAI stats
         # print(f"Model name: {response.llm_output.get('model_name', '')}")

environment.yaml

@@ -18,3 +18,4 @@ dependencies:
     - streamlit==1.24.1
     - streamlit-folium==0.12.0
     - watchdog==3.0.0
+    - aim==3.17.5

tools/geopy/geocode.py

@@ -21,6 +21,8 @@ class GeopyGeocodeTool(BaseTool):
     def _run(self, place: str) -> tuple:
         locator = Nominatim(user_agent="geocode")
         location = locator.geocode(place)
+        if location is None:
+            return ("geocode", "Not a recognised address in Nomatim.")
         return ("geocode", (location.latitude, location.longitude))
 
     def _arun(self, place: str):

tools/osmnx/geometry.py

@@ -7,19 +7,19 @@
 
 
 class PlaceWithTags(BaseModel):
-    "Name of a place and tags in OSM."
+    "Name of a place on the map and tags in OSM."
 
-    place: str = Field(..., description="name of a place")
-    tags: Dict[str, str] = Field(..., description="open street maps tags")
+    place: str = Field(..., description="name of a place on the map.")
+    tags: Dict[str, str] = Field(..., description="open street maps tags.")
 
 
 class OSMnxGeometryTool(BaseTool):
     """Tool to query geometries from Open Street Map (OSM)."""
 
     name: str = "geometry"
     args_schema: Type[BaseModel] = PlaceWithTags
-    description: str = "Use this tool to get geometry of different features of a place like building footprints, parks, lakes, hospitals, schools etc. \
-    Pass the name of the place & relevant tags of OSM as args."
+    description: str = "Use this tool to get geometry of different features of the place like building footprints, parks, lakes, hospitals, schools etc. \
+    Pass the name of the place & tags of OSM as args."
     return_direct = True
 
     def _run(self, place: str, tags: Dict[str, str]) -> gpd.GeoDataFrame: