mirror of
https://github.com/microsoft/semantic-kernel.git
synced 2026-10-12 02:19:06 +00:00
Python: feat(oracle): add new Oracle connector for Semantic Kernel (#13229)
…ync support ### Motivation and Context <!-- Thank you for your contribution to the semantic-kernel repo! Please help reviewers and future users, providing the following information: 1. Why is this change required? 2. What problem does it solve? 3. What scenario does it contribute to? 4. If it fixes an open issue, please link to the issue here. --> This change is required to enable Semantic Kernel users to store and retrieve embeddings using Oracle databases. Currently, Semantic Kernel supports vector storage for several backends, but Oracle was missing. This connector solves that gap by providing full async support, native VECTOR type handling, and vector index management. ### Description <!-- Describe your changes, the overall approach, the underlying design. These notes will help understanding how your code works. Thanks! --> This PR introduces a new Oracle connector for Semantic Kernel with the following features: - Asynchronous upsert, get, delete and search operations for memory records. - Native Oracle VECTOR type support for storing embeddings efficiently. - Support for HNSW and IVFFLAT vector indexes for similarity search. - Integration with Semantic Kernel collections, enabling semantic search and memory operations. - Comprehensive unit tests to ensure correctness and stability. The connector is designed to work seamlessly with existing Semantic Kernel memory abstractions and follows the same async patterns as other vector stores. Integration tests have also been implemented and verified locally; however, they are not included in this PR because the current CI environment setup for Oracle Database support is unknown. Once guidance is provided on Oracle DB availability in the CI pipeline, integration tests can be enabled and added in a follow-up PR. ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [x] The code builds clean without any errors or warnings - [x] 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 - [x] All unit tests pass, and I have added new tests where possible - [x] I didn't break anyone 😄
This commit is contained in:
@@ -125,6 +125,9 @@ onnx = [
|
||||
"onnxruntime==1.22.1",
|
||||
"onnxruntime-genai==0.9.0"
|
||||
]
|
||||
oracledb = [
|
||||
"oracledb >= 3.4.1"
|
||||
]
|
||||
pandas = [
|
||||
"pandas ~= 2.2"
|
||||
]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1316,7 +1316,9 @@ class VectorStoreCollection(VectorStoreRecordHandler[TKey, TModel], Generic[TKey
|
||||
try:
|
||||
results = await self._inner_upsert(data if isinstance(data, list) else [data], **kwargs) # type: ignore
|
||||
except Exception as exc:
|
||||
raise VectorStoreOperationException(f"Error upserting record(s): {exc}") from exc
|
||||
raise VectorStoreOperationException(
|
||||
f"Error upserting record(s) into collection '{self.collection_name}': {exc}"
|
||||
) from exc
|
||||
if batch or self._container_mode:
|
||||
return results
|
||||
return results[0]
|
||||
|
||||
@@ -0,0 +1,421 @@
|
||||
# Copyright (c) 2025, Oracle Corporation. All rights reserved. # noqa: CPY001
|
||||
|
||||
from array import array
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from typing import Annotated
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import oracledb
|
||||
import pandas as pd
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from semantic_kernel.connectors.oracle import OracleCollection, OracleStore
|
||||
from semantic_kernel.data.vector import (
|
||||
DistanceFunction,
|
||||
IndexKind,
|
||||
VectorStoreCollectionDefinition,
|
||||
VectorStoreField,
|
||||
vectorstoremodel,
|
||||
)
|
||||
|
||||
|
||||
@vectorstoremodel
|
||||
@dataclass
|
||||
class SimpleModel:
|
||||
id: Annotated[int, VectorStoreField("key")]
|
||||
vector: Annotated[
|
||||
list[float] | None,
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
type="float",
|
||||
dimensions=3,
|
||||
index_kind=IndexKind.HNSW,
|
||||
distance_function=DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE,
|
||||
),
|
||||
] = None
|
||||
|
||||
|
||||
def PandasDataframeModel(record) -> tuple:
|
||||
definition = VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id", type="int"),
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
name="embedding",
|
||||
type="float32",
|
||||
dimensions=5,
|
||||
distance_function=DistanceFunction.COSINE_DISTANCE,
|
||||
index_kind=IndexKind.IVF_FLAT,
|
||||
),
|
||||
],
|
||||
to_dict=lambda record, **_: record.to_dict(orient="records"),
|
||||
from_dict=lambda records, **_: pd.DataFrame(records),
|
||||
container_mode=True,
|
||||
)
|
||||
df = pd.DataFrame([record]) if isinstance(record, dict) else pd.DataFrame(record)
|
||||
return definition, df
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def mock_connection_pool():
|
||||
mock_conn = AsyncMock()
|
||||
mock_conn.fetchall = AsyncMock(return_value=[("COLL1",), ("COLL2",)])
|
||||
|
||||
mock_context_manager = AsyncMock()
|
||||
mock_context_manager.__aenter__.return_value = mock_conn
|
||||
mock_context_manager.__aexit__.return_value = None
|
||||
|
||||
pool = MagicMock(spec=oracledb.AsyncConnectionPool)
|
||||
pool.acquire.return_value = mock_context_manager
|
||||
|
||||
return pool
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def oracle_store(mock_connection_pool):
|
||||
return OracleStore(
|
||||
connection_pool=mock_connection_pool,
|
||||
db_schema="MY_SCHEMA",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_collection_names_with_schema(oracle_store):
|
||||
names = await oracle_store.list_collection_names()
|
||||
assert names == ["COLL1", "COLL2"]
|
||||
|
||||
|
||||
def test_get_collection_returns_oracle_collection(oracle_store):
|
||||
collection = oracle_store.get_collection(
|
||||
SimpleModel,
|
||||
collection_name="TEST",
|
||||
)
|
||||
assert isinstance(collection, OracleCollection)
|
||||
assert collection.collection_name == "TEST"
|
||||
assert collection.db_schema == "MY_SCHEMA"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_collection_exists_true(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.fetchone = AsyncMock(return_value=(1,))
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
SimpleModel,
|
||||
collection_name="EXISTING",
|
||||
)
|
||||
result = await collection.collection_exists()
|
||||
assert result is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_collection_exists_false(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.fetchone = AsyncMock(return_value=None)
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
SimpleModel,
|
||||
collection_name="MISSING",
|
||||
)
|
||||
result = await collection.collection_exists()
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_collection_exists_creates_when_missing(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.fetchone = AsyncMock(return_value=False)
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="NEW")
|
||||
await collection.ensure_collection_exists()
|
||||
conn.execute.assert_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_table_with_get_collection(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="MY_COLLECTION")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
mock_conn.execute.assert_awaited()
|
||||
sql_statements = [args[0].lower() for name, args, _ in mock_conn.mock_calls if name == "execute"]
|
||||
|
||||
assert any("create table" in sql for sql in sql_statements)
|
||||
assert any("my_collection" in sql for sql in sql_statements)
|
||||
assert any("vector(3 , float64)" in sql for sql in sql_statements)
|
||||
|
||||
mock_conn.commit.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pandasDataframe_with_get_collection(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
records = {
|
||||
"id": 1,
|
||||
"embedding": [1.1, 2.2, 3.3],
|
||||
}
|
||||
|
||||
definition, _ = PandasDataframeModel(records)
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
collection_name="MY_COLLECTION",
|
||||
record_type=pd.DataFrame,
|
||||
definition=definition,
|
||||
)
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
mock_conn.execute.assert_awaited()
|
||||
sql_statements = [args[0].lower() for name, args, _ in mock_conn.mock_calls if name == "execute"]
|
||||
|
||||
assert any("create table" in sql for sql in sql_statements)
|
||||
assert any("my_collection" in sql for sql in sql_statements)
|
||||
assert any("vector(5 , float32)" in sql for sql in sql_statements)
|
||||
|
||||
mock_conn.commit.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_index_creation_distance_with_get_collection(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="COLLECTION_WITH_INDEX")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
mock_conn.execute.assert_awaited()
|
||||
called_sql = mock_conn.execute.call_args[0][0].lower()
|
||||
|
||||
assert "create vector index" in called_sql
|
||||
assert "collection_with_index_vector_idx" in called_sql
|
||||
assert "inmemory neighbor graph" in called_sql
|
||||
assert "distance euclidean_squared" in called_sql
|
||||
|
||||
mock_conn.commit.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_collection_deleted(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.fetchone = AsyncMock(return_value=(1,))
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="TO_DELETE")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.ensure_collection_deleted()
|
||||
|
||||
assert any("DROP TABLE" in str(call.args[0]) for call in conn.execute.call_args_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="MY_COLLECTION")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.upsert(SimpleModel(id=1, vector=[0.1, 0.2, 0.3]))
|
||||
|
||||
mock_conn.executemany.assert_called_once()
|
||||
|
||||
merge_sql, params = mock_conn.executemany.call_args[0]
|
||||
|
||||
assert merge_sql.startswith('MERGE INTO "MY_SCHEMA"."MY_COLLECTION"')
|
||||
assert 'UPDATE SET t."vector"' in merge_sql
|
||||
assert "WHEN NOT MATCHED THEN" in merge_sql
|
||||
assert 'INSERT ("id", "vector")' in merge_sql
|
||||
|
||||
expected_param = (1, array("d", [0.1, 0.2, 0.3]))
|
||||
assert params[0] == expected_param
|
||||
assert mock_conn.commit.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_with_include_vectors(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_conn.description = [("id",), ("vector",)]
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="MY_COLLECTION")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.upsert(SimpleModel(id=1, vector=[0.1, 0.2, 0.3]))
|
||||
|
||||
mock_conn.fetchall.return_value = [(1, [0.1, 0.2, 0.3])]
|
||||
results = await collection.get([1], include_vectors=True)
|
||||
assert results.id == 1
|
||||
assert results.vector == [0.1, 0.2, 0.3]
|
||||
|
||||
executed_sql = [args[0] for args, _ in mock_conn.fetchall.call_args_list]
|
||||
assert any('SELECT "id" AS "id", "vector" AS "vector"' in sql for sql in executed_sql), (
|
||||
"Expected vector column to be selected when include_vectors=True"
|
||||
)
|
||||
|
||||
mock_conn.fetchall.reset_mock()
|
||||
mock_conn.fetchall.return_value = [(1, None)]
|
||||
results = await collection.get([1], include_vectors=False)
|
||||
assert results.id == 1
|
||||
assert results.vector is None
|
||||
|
||||
executed_sql = [args[0] for args, _ in mock_conn.fetchall.call_args_list]
|
||||
assert any('SELECT "id" AS "id", "vector" AS "vector"' not in sql for sql in executed_sql), (
|
||||
"Vector column should not be selected when include_vectors=False"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_and_get(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.description = [("id",), ("vector",)]
|
||||
conn.fetchall = AsyncMock(return_value=[(1, [0.1, 0.2, 0.3])])
|
||||
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="TEST")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.upsert(SimpleModel(id=1, vector=[0.1, 0.2, 0.3]))
|
||||
assert conn.executemany.await_count >= 1 or conn.execute.await_count >= 1, (
|
||||
"Expected upsert to call executemany() or execute()"
|
||||
)
|
||||
|
||||
results = await collection.get([1], include_vectors=True)
|
||||
|
||||
assert results.id == 1
|
||||
assert results.vector == [0.1, 0.2, 0.3]
|
||||
conn.fetchall.assert_awaited_once()
|
||||
conn.fetchall.reset_mock()
|
||||
conn.fetchall.return_value = [(1, None)]
|
||||
conn.description = [("id",)]
|
||||
|
||||
results = await collection.get([1], include_vectors=False)
|
||||
assert results.id == 1
|
||||
assert results.vector is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_record(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
SimpleModel,
|
||||
collection_name="MY_COLLECTION",
|
||||
)
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.delete([1])
|
||||
conn.executemany.assert_awaited()
|
||||
called_sql = conn.executemany.call_args[0][0].lower()
|
||||
assert "delete" in called_sql
|
||||
assert "my_collection" in called_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search(oracle_store, mock_connection_pool):
|
||||
class MockCursor:
|
||||
def __init__(self):
|
||||
self.execute_called_with = []
|
||||
self._rows = [(1, [0.1, 0.2, 0.3])]
|
||||
self.description = [SimpleNamespace(name="id"), SimpleNamespace(name="vector")]
|
||||
self._i = 0
|
||||
|
||||
async def execute(self, sql, binds=None):
|
||||
self.execute_called_with.append((sql, binds))
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
return None
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
return None
|
||||
|
||||
def __aiter__(self):
|
||||
self._i = 0
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._i >= len(self._rows):
|
||||
raise StopAsyncIteration
|
||||
|
||||
r = self._rows[self._i]
|
||||
self._i += 1
|
||||
return r
|
||||
|
||||
mock_cursor = MockCursor()
|
||||
|
||||
class MockConnection:
|
||||
def __init__(self, cur):
|
||||
self._cur = cur
|
||||
self.inputtypehandler = None
|
||||
self.outputtypehandler = None
|
||||
self.execute = AsyncMock()
|
||||
self.commit = AsyncMock()
|
||||
|
||||
def cursor(self):
|
||||
return self._cur
|
||||
|
||||
mock_conn = MockConnection(mock_cursor)
|
||||
|
||||
class MockAcquire:
|
||||
def __init__(self, conn):
|
||||
self._conn = conn
|
||||
|
||||
def __await__(self):
|
||||
async def _():
|
||||
return self
|
||||
|
||||
return _().__await__()
|
||||
|
||||
async def __aenter__(self):
|
||||
return self._conn
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
return None
|
||||
|
||||
mock_connection_pool.acquire = lambda **kwargs: MockAcquire(mock_conn)
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
model=SimpleModel, record_type=SimpleModel, collection_name="MY_COLLECTION"
|
||||
)
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
assert mock_conn.execute.await_count >= 1
|
||||
|
||||
ks_results = await collection.search(
|
||||
vector_property_name="vector",
|
||||
vector=[0.1, 0.2, 0.3],
|
||||
top=1,
|
||||
filter=lambda x: x.id in [1, 7, 9],
|
||||
include_vectors=True,
|
||||
)
|
||||
results = [r async for r in ks_results.results]
|
||||
|
||||
assert mock_cursor.execute_called_with, "cursor.execute was not called"
|
||||
sql, binds = mock_cursor.execute_called_with[0]
|
||||
assert "SELECT" in sql.upper()
|
||||
assert '"id", "vector", VECTOR_DISTANCE' in sql
|
||||
expected_where = 'WHERE "id" IN (:bind_val1, :bind_val2, :bind_val3)'
|
||||
assert expected_where in sql
|
||||
|
||||
assert binds is None or isinstance(binds[0], array)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].record.id == 1
|
||||
assert results[0].record.vector == [0.1, 0.2, 0.3]
|
||||
Reference in New Issue
Block a user