From 1545dc4d872f4b93b472cc515da04f9e62c6243c Mon Sep 17 00:00:00 2001
From: Sid Murching <sid.murching@databricks.com>
Date: Tue, 3 Dec 2024 17:06:47 -0800
Subject: [PATCH] WIP

Signed-off-by: Sid Murching <sid.murching@databricks.com>
---
 .../langchain/src/databricks_langchain/vector_search.py   | 8 ++++++++
 src/databricks_ai_bridge/vector_search.py                 | 2 ++
 2 files changed, 10 insertions(+)
 create mode 100644 integrations/langchain/src/databricks_langchain/vector_search.py
 create mode 100644 src/databricks_ai_bridge/vector_search.py

diff --git a/integrations/langchain/src/databricks_langchain/vector_search.py b/integrations/langchain/src/databricks_langchain/vector_search.py
new file mode 100644
index 0000000..728f263
--- /dev/null
+++ b/integrations/langchain/src/databricks_langchain/vector_search.py
@@ -0,0 +1,8 @@
+from databricks_langchain.vector_search import DatabricksVectorSearch
+class VectorSearchRetrieverTool():
+    def __init__(self, *args, **kwargs):
+        vector_store = DatabricksVectorSearch(
+            endpoint=endpoint_name,
+            index_name=index_name,
+        )
+        return vector_store.as_retriever().as_tool()
diff --git a/src/databricks_ai_bridge/vector_search.py b/src/databricks_ai_bridge/vector_search.py
new file mode 100644
index 0000000..211830b
--- /dev/null
+++ b/src/databricks_ai_bridge/vector_search.py
@@ -0,0 +1,2 @@
+# Put common logic for exposing Databricks vector search as a tool here
+# for example: querying a VS endpoint, getting metadata about the index, ...