Create connection.py

connection.py

@@ -0,0 +1,49 @@
+import pandas as pd
+import pinecone
+import streamlit as st
+from streamlit.connections import ExperimentalBaseConnection
+
+
+class PineconeConnection(ExperimentalBaseConnection):
+    def __init__(
+        self,
+        connection_name: str,
+        environment=None,
+        api_key=None,
+        **kwargs,
+    ) -> None:
+        self.environment = environment
+        self.api_key = api_key
+        super().__init__(connection_name, **kwargs)
+
+    def _connect(self):
+        api_key = self.api_key or self._secrets.get("Pinecone_API_KEY")
+        environment = self.environment
+        return pinecone.init(api_key=api_key, environment=environment)
+
+    def list_indexes(self):
+        self._connect()
+        self.indexes = pinecone.list_indexes()
+        return self.indexes
+
+    def _connect_index(self, index_name):
+        self._connect()
+        self.index_name = index_name
+        self.index = pinecone.Index(index_name)
+        return self.index
+
+    def query(
+        self, index_name: str, query_vector, top_k: int = 5, ttl: int = 3600, **kwargs
+    ) -> dict:
+        @st.cache_resource(ttl=ttl)
+        def _query(index_name: str, query_vector, top_k: int = 5, **kwargs):
+            index = self._connect_index(index_name)
+            query_results = index.query(query_vector, top_k=top_k, **kwargs)
+            results = list(query_results["matches"])
+            return results
+
+        results = _query(index_name, query_vector, top_k, **kwargs)
+        return results
+
+    def cursor(self):
+        return self._connect()