Fix load_crash_data retriever function
#4
by tigerstride - opened
- 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 |
-
|
|
|
|
| 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 |
|