unit and integration tests for chat message histories and vector stores

libs/arangodb/tests/integration_tests/chat_message_histories/test_arangodb.py

@@ -1,8 +1,14 @@
+import os
+
 import pytest
+from arango import ArangoClient
 from arango.database import StandardDatabase
+from arango.exceptions import ArangoError
 from langchain_core.messages import AIMessage, HumanMessage
 
 from langchain_arangodb.chat_message_histories.arangodb import ArangoChatMessageHistory
+from langchain_arangodb.graphs.arangodb_graph import ArangoGraph
+from tests.integration_tests.utils import ArangoCredentials
 
 
 @pytest.mark.usefixtures("clear_arangodb_database")
@@ -44,3 +50,108 @@ def test_add_messages(db: StandardDatabase) -> None:
     message_store_another.clear()
     assert len(message_store.messages) == 0
     assert len(message_store_another.messages) == 0
+
+
+@pytest.mark.usefixtures("clear_arangodb_database")
+def test_add_messages_graph_object(arangodb_credentials: ArangoCredentials) -> None:
+    """Basic testing: Passing driver through graph object."""
+    graph = ArangoGraph.from_db_credentials(
+        url=arangodb_credentials["url"],
+        username=arangodb_credentials["username"],
+        password=arangodb_credentials["password"],
+    )
+
+    # rewrite env for testing
+    old_username = os.environ.get("ARANGO_USERNAME", "root")
+    os.environ["ARANGO_USERNAME"] = "foo"
+
+    message_store = ArangoChatMessageHistory("23334", db=graph.db)
+    message_store.clear()
+    assert len(message_store.messages) == 0
+    message_store.add_user_message("Hello! Language Chain!")
+    message_store.add_ai_message("Hi Guys!")
+    # Now check if the messages are stored in the database correctly
+    assert len(message_store.messages) == 2
+
+    # Restore original environment
+    os.environ["ARANGO_USERNAME"] = old_username
+
+
+def test_invalid_credentials(arangodb_credentials: ArangoCredentials) -> None:
+    """Test initializing with invalid credentials raises an authentication error."""
+    with pytest.raises(ArangoError) as exc_info:
+        client = ArangoClient(arangodb_credentials["url"])
+        db = client.db(username="invalid_username", password="invalid_password")
+        # Try to perform a database operation to trigger an authentication error
+        db.collections()
+
+    # Check for any authentication-related error message
+    error_msg = str(exc_info.value)
+    # Just check for "error" which should be in any auth error
+    assert "not authorized" in error_msg
+
+
+@pytest.mark.usefixtures("clear_arangodb_database")
+def test_arangodb_message_history_clear_messages(
+    db: StandardDatabase,
+) -> None:
+    """Test adding multiple messages at once to ArangoChatMessageHistory."""
+    # Specify a custom collection name that includes the session_id
+    collection_name = "chat_history_123"
+    message_history = ArangoChatMessageHistory(
+        session_id="123", db=db, collection_name=collection_name
+    )
+    message_history.add_messages(
+        [
+            HumanMessage(content="You are a helpful assistant."),
+            AIMessage(content="Hello"),
+        ]
+    )
+    assert len(message_history.messages) == 2
+    assert isinstance(message_history.messages[0], HumanMessage)
+    assert isinstance(message_history.messages[1], AIMessage)
+    assert message_history.messages[0].content == "You are a helpful assistant."
+    assert message_history.messages[1].content == "Hello"
+
+    message_history.clear()
+    assert len(message_history.messages) == 0
+
+    # Verify all messages are removed but collection still exists
+    assert db.has_collection(message_history._collection_name)
+    assert message_history._collection_name == collection_name
+
+
+@pytest.mark.usefixtures("clear_arangodb_database")
+def test_arangodb_message_history_clear_session_collection(
+    db: StandardDatabase,
+) -> None:
+    """Test clearing messages and removing the collection for a session."""
+    # Create a test collection specific to the session
+    session_id = "456"
+    collection_name = f"chat_history_{session_id}"
+
+    if not db.has_collection(collection_name):
+        db.create_collection(collection_name)
+
+    message_history = ArangoChatMessageHistory(
+        session_id=session_id, db=db, collection_name=collection_name
+    )
+
+    message_history.add_messages(
+        [
+            HumanMessage(content="You are a helpful assistant."),
+            AIMessage(content="Hello"),
+        ]
+    )
+    assert len(message_history.messages) == 2
+
+    # Clear messages
+    message_history.clear()
+    assert len(message_history.messages) == 0
+
+    # The collection should still exist after clearing messages
+    assert db.has_collection(collection_name)
+
+    # Delete the collection (equivalent to delete_session_node in Neo4j)
+    db.delete_collection(collection_name)
+    assert not db.has_collection(collection_name)

libs/arangodb/tests/integration_tests/vectorstores/fake_embeddings.py

@@ -0,0 +1,111 @@
+"""Fake Embedding class for testing purposes."""
+
+import math
+from typing import List
+
+from langchain_core.embeddings import Embeddings
+
+fake_texts = ["foo", "bar", "baz"]
+
+
+class FakeEmbeddings(Embeddings):
+    """Fake embeddings functionality for testing."""
+
+    def __init__(self, dimension: int = 10):
+        if dimension < 1:
+            raise ValueError(
+                "Dimension must be at least 1 for this FakeEmbeddings style."
+            )
+        self.dimension = dimension
+        # global_fake_texts maps query texts to the 'i' in [1.0]*(dim-1) + [float(i)]
+        self.global_fake_texts = ["foo", "bar", "baz", "qux", "quux", "corge", "grault"]
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Return simple embeddings.
+        Embeddings encode each text as its index."""
+        if self.dimension == 1:
+            # Special case for dimension 1: just use the index
+            return [[float(i)] for i in range(len(texts))]
+        else:
+            return [
+                [1.0] * (self.dimension - 1) + [float(i)] for i in range(len(texts))
+            ]
+
+    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
+        return self.embed_documents(texts)
+
+    def embed_query(self, text: str) -> List[float]:
+        """Return constant query embeddings.
+        Embeddings are identical to embed_documents(texts)[0].
+        Distance to each text will be that text's index,
+        as it was passed to embed_documents."""
+        try:
+            idx = self.global_fake_texts.index(text)
+            val = float(idx)
+        except ValueError:
+            # Text not in global_fake_texts, use a default 'unknown query' value
+            val = -1.0
+
+        if self.dimension == 1:
+            return [val]  # Corrected: List[float]
+        else:
+            return [1.0] * (self.dimension - 1) + [val]
+
+    async def aembed_query(self, text: str) -> List[float]:
+        return self.embed_query(text)
+
+    @property
+    def identifer(self) -> str:
+        return "fake"
+
+
+class ConsistentFakeEmbeddings(FakeEmbeddings):
+    """Fake embeddings which remember all the texts seen so far to return consistent
+    vectors for the same texts."""
+
+    def __init__(self, dimensionality: int = 10) -> None:
+        self.known_texts: List[str] = []
+        self.dimensionality = dimensionality
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Return consistent embeddings for each text seen so far."""
+        out_vectors = []
+        for text in texts:
+            if text not in self.known_texts:
+                self.known_texts.append(text)
+            vector = [float(1.0)] * (self.dimensionality - 1) + [
+                float(self.known_texts.index(text))
+            ]
+            out_vectors.append(vector)
+        return out_vectors
+
+    def embed_query(self, text: str) -> List[float]:
+        """Return consistent embeddings for the text, if seen before, or a constant
+        one if the text is unknown."""
+        return self.embed_documents([text])[0]
+
+
+class AngularTwoDimensionalEmbeddings(Embeddings):
+    """
+    From angles (as strings in units of pi) to unit embedding vectors on a circle.
+    """
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        Make a list of texts into a list of embedding vectors.
+        """
+        return [self.embed_query(text) for text in texts]
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Convert input text to a 'vector' (list of floats).
+        If the text is a number, use it as the angle for the
+        unit vector in units of pi.
+        Any other input text becomes the singular result [0, 0] !
+        """
+        try:
+            angle = float(text)
+            return [math.cos(angle * math.pi), math.sin(angle * math.pi)]
+        except ValueError:
+            # Assume: just test string, no attention is paid to values.
+            return [0.0, 0.0]