Skip to content

Commit 9dbbf56

Browse files
DevTitan-AIOverlordjsbilgianomnaco
authored
Python: Add astra memory (microsoft#4374)
### Motivation and Context We are going to add [Astra](https://docs.datastax.com/en/home/docs/index.html) as a memory store. DataStax Astra DB is a serverless vector database that’s perfect for managing mission-critical AI workloads. It’s built on Apache Cassandra®, making Astra DB a highly scalable, reliable database technology. It’s a powerful all-in-one data storage solution, ideal for Generative AI projects. ### Description - Add astra into python/semantic_kernel/connectors/memory/astradb - Add test into python/tests/integration/connectors/memory/test_astradb.py ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [ ] The code builds clean without any errors or warnings - [ ] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [ ] All unit tests pass, and I have added new tests where possible - [ ] I didn't break anyone 😄 --------- Co-authored-by: jsbilgi <[email protected]> Co-authored-by: Obioma Anomnachi <[email protected]>
1 parent 11cf8ea commit 9dbbf56

8 files changed

Lines changed: 829 additions & 0 deletions

File tree

‎python/.env.example‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,3 +18,7 @@ AZCOSMOS_API = "" // should be mongo-vcore for now, as CosmosDB only supports ve
1818
AZCOSMOS_CONNSTR = ""
1919
AZCOSMOS_DATABASE_NAME = ""
2020
AZCOSMOS_CONTAINER_NAME = ""
21+
ASTRADB_APP_TOKEN="" // Starts with AstraCS:
22+
ASTRADB_ID=""
23+
ASTRADB_REGION=""
24+
ASTRADB_KEYSPACE=""

‎python/semantic_kernel/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from semantic_kernel.utils.logging import setup_logging
1717
from semantic_kernel.utils.null_logger import NullLogger
1818
from semantic_kernel.utils.settings import (
19+
astradb_settings_from_dot_env,
1920
azure_aisearch_settings_from_dot_env,
2021
azure_aisearch_settings_from_dot_env_as_dict,
2122
azure_cosmos_db_settings_from_dot_env,
@@ -39,6 +40,7 @@
3940
"azure_aisearch_settings_from_dot_env_as_dict",
4041
"postgres_settings_from_dot_env",
4142
"pinecone_settings_from_dot_env",
43+
"astradb_settings_from_dot_env",
4244
"bing_search_settings_from_dot_env",
4345
"mongodb_atlas_settings_from_dot_env",
4446
"google_palm_settings_from_dot_env",

‎python/semantic_kernel/connectors/memory/astradb/__init__.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
# Copyright (c) Microsoft. All rights reserved.
2+
3+
from semantic_kernel.connectors.memory.astradb.astradb_memory_store import (
4+
AstraDBMemoryStore,
5+
)
6+
7+
__all__ = ["AstraDBMemoryStore"]

‎python/semantic_kernel/connectors/memory/astradb/astra_client.py‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
import json
2+
from typing import Dict, List, Optional
3+
4+
import aiohttp
5+
6+
from semantic_kernel.connectors.memory.astradb.utils import AsyncSession
7+
8+
9+
class AstraClient:
10+
def __init__(
11+
self,
12+
astra_id: str,
13+
astra_region: str,
14+
astra_application_token: str,
15+
keyspace_name: str,
16+
embedding_dim: int,
17+
similarity_function: str,
18+
session: Optional[aiohttp.ClientSession] = None,
19+
):
20+
self.astra_id = astra_id
21+
self.astra_application_token = astra_application_token
22+
self.astra_region = astra_region
23+
self.keyspace_name = keyspace_name
24+
self.embedding_dim = embedding_dim
25+
self.similarity_function = similarity_function
26+
27+
self.request_base_url = (
28+
f"https://{self.astra_id}-{self.astra_region}.apps.astra.datastax.com/api/json/v1/{self.keyspace_name}"
29+
)
30+
self.request_header = {
31+
"x-cassandra-token": self.astra_application_token,
32+
"Content-Type": "application/json",
33+
}
34+
self._session = session
35+
36+
async def _run_query(self, request_url: str, query: Dict):
37+
async with AsyncSession(self._session) as session:
38+
async with session.post(request_url, data=json.dumps(query), headers=self.request_header) as response:
39+
if response.status == 200:
40+
response_dict = await response.json()
41+
if "errors" in response_dict:
42+
raise Exception(f"Astra DB request error - {response_dict['errors']}")
43+
else:
44+
return response_dict
45+
else:
46+
raise Exception(f"Astra DB not available. Status : {response}")
47+
48+
async def find_collections(self, include_detail: bool = True):
49+
query = {"findCollections": {"options": {"explain": include_detail}}}
50+
result = await self._run_query(self.request_base_url, query)
51+
return result["status"]["collections"]
52+
53+
async def find_collection(self, collection_name: str):
54+
collections = await self.find_collections(False)
55+
found = False
56+
for collection in collections:
57+
if collection == collection_name:
58+
found = True
59+
break
60+
return found
61+
62+
async def create_collection(
63+
self,
64+
collection_name: str,
65+
embedding_dim: Optional[int] = None,
66+
similarity_function: Optional[str] = None,
67+
):
68+
query = {
69+
"createCollection": {
70+
"name": collection_name,
71+
"options": {
72+
"vector": {
73+
"dimension": embedding_dim if embedding_dim is not None else self.embedding_dim,
74+
"metric": similarity_function if similarity_function is not None else self.similarity_function,
75+
}
76+
},
77+
}
78+
}
79+
result = await self._run_query(self.request_base_url, query)
80+
return True if result["status"]["ok"] == 1 else False
81+
82+
async def delete_collection(self, collection_name: str):
83+
query = {"deleteCollection": {"name": collection_name}}
84+
result = await self._run_query(self.request_base_url, query)
85+
return True if result["status"]["ok"] == 1 else False
86+
87+
def _build_request_collection_url(self, collection_name: str):
88+
return f"{self.request_base_url}/{collection_name}"
89+
90+
async def find_documents(
91+
self,
92+
collection_name: str,
93+
filter: Optional[Dict] = None,
94+
vector: Optional[List[float]] = None,
95+
limit: Optional[int] = None,
96+
include_vector: Optional[bool] = None,
97+
include_similarity: Optional[bool] = None,
98+
) -> List[Dict]:
99+
find_query = {}
100+
101+
if filter is not None:
102+
find_query["filter"] = filter
103+
104+
if vector is not None:
105+
find_query["sort"] = {"$vector": vector}
106+
107+
if include_vector is not None and include_vector is False:
108+
find_query["projection"] = {"$vector": 0}
109+
110+
if limit is not None:
111+
find_query["options"] = {"limit": limit}
112+
113+
if include_similarity is not None:
114+
if "options" in find_query:
115+
find_query["options"]["includeSimilarity"] = int(include_similarity)
116+
else:
117+
find_query["options"] = {"includeSimilarity": int(include_similarity)}
118+
119+
query = {"find": find_query}
120+
result = await self._run_query(self._build_request_collection_url(collection_name), query)
121+
return result["data"]["documents"]
122+
123+
async def insert_document(self, collection_name: str, document: Dict) -> str:
124+
query = {"insertOne": {"document": document}}
125+
result = await self._run_query(self._build_request_collection_url(collection_name), query)
126+
return result["status"]["insertedIds"][0]
127+
128+
async def insert_documents(self, collection_name: str, documents: List[Dict]) -> List[str]:
129+
query = {"insertMany": {"documents": documents}}
130+
result = await self._run_query(self._build_request_collection_url(collection_name), query)
131+
return result["status"]["insertedIds"]
132+
133+
async def update_document(self, collection_name: str, filter: Dict, update: Dict, upsert: bool = True) -> Dict:
134+
query = {
135+
"findOneAndUpdate": {
136+
"filter": filter,
137+
"update": update,
138+
"options": {"returnDocument": "after", "upsert": upsert},
139+
}
140+
}
141+
result = await self._run_query(self._build_request_collection_url(collection_name), query)
142+
return result["status"]
143+
144+
async def update_documents(self, collection_name: str, filter: Dict, update: Dict):
145+
query = {
146+
"updateMany": {
147+
"filter": filter,
148+
"update": update,
149+
}
150+
}
151+
result = await self._run_query(self._build_request_collection_url(collection_name), query)
152+
return result["status"]
153+
154+
async def delete_documents(self, collection_name: str, filter: Dict) -> int:
155+
query = {"deleteMany": {"filter": filter}}
156+
result = await self._run_query(self._build_request_collection_url(collection_name), query)
157+
return result["status"]["deletedCount"]

0 commit comments

Comments
 (0)