fix databricksrm (#1967)

stanfordnlp · Dec 20, 2024 · dbc96d4 · dbc96d4
1 parent 6085098
commit dbc96d4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dspy/retrieve/databricks_rm.py b/dspy/retrieve/databricks_rm.py
@@ -288,7 +288,7 @@ def forward(
                     page_content=doc[self.text_column_name],
                     metadata={
                         "doc_id": self._extract_doc_ids(doc),
-                        "doc_uri": doc[self.docs_uri_column_name],
+                        "doc_uri": doc[self.docs_uri_column_name] if self.docs_uri_column_name else None,
                     }
                     | self._get_extra_columns(doc),
                     type="Document",
@@ -300,7 +300,7 @@ def forward(
             return Prediction(
                 docs=[doc[self.text_column_name] for doc in sorted_docs],
                 doc_ids=[self._extract_doc_ids(doc) for doc in sorted_docs],
-                doc_uris=[doc[self.docs_uri_column_name] for doc in sorted_docs],
+                doc_uris=[doc[self.docs_uri_column_name] for doc in sorted_docs] if self.docs_uri_column_name else None,
                 extra_columns=[self._get_extra_columns(item) for item in sorted_docs],
             )