tigerstride commited on
Commit
b80a29a
·
verified ·
1 Parent(s): 8aec252

Fix load_crash_data retriever function (#4)

Browse files

- Fix load_crash_data retriever function (15c3175f9f059cae97df0146478f9d7ad043be2c)

Files changed (1) hide show
  1. retriever.py +4 -3
retriever.py CHANGED
@@ -1,6 +1,7 @@
1
  from smolagents import Tool
2
  from langchain_community.retrievers import BM25Retriever
3
  from langchain.docstore.document import Document
 
4
  import datasets
5
  import pandas as pd
6
  import requests
@@ -50,7 +51,6 @@ def load_guest_dataset():
50
  # Return the tool
51
  return GuestInfoRetrieverTool(docs)
52
 
53
-
54
  def load_crash_data():
55
  # Public Philly crash data CSV (small demo file)
56
  url = "https://phl.carto.com/api/v2/sql?q=SELECT+*+FROM+vision_zero_crash+LIMIT+5000&format=csv"
@@ -72,6 +72,7 @@ def load_crash_data():
72
  f"Weather: {row.get('weather', 'N/A')}"
73
  ])
74
  docs.append(Document(page_content=content, metadata={"id": row.get("id", "unknown")}))
75
-
76
- return docs
 
77
 
 
1
  from smolagents import Tool
2
  from langchain_community.retrievers import BM25Retriever
3
  from langchain.docstore.document import Document
4
+ from tools import CrashInfoRetrieverTool # Import the CrashInfoRetrieverTool
5
  import datasets
6
  import pandas as pd
7
  import requests
 
51
  # Return the tool
52
  return GuestInfoRetrieverTool(docs)
53
 
 
54
  def load_crash_data():
55
  # Public Philly crash data CSV (small demo file)
56
  url = "https://phl.carto.com/api/v2/sql?q=SELECT+*+FROM+vision_zero_crash+LIMIT+5000&format=csv"
 
72
  f"Weather: {row.get('weather', 'N/A')}"
73
  ])
74
  docs.append(Document(page_content=content, metadata={"id": row.get("id", "unknown")}))
75
+
76
+ # Return an instance of CrashInfoRetrieverTool
77
+ return CrashInfoRetrieverTool(docs)
78