diff --git a/examples/basic/vector_ollama.py b/examples/basic/vector_ollama.py new file mode 100644 index 0000000..c7b2399 --- /dev/null +++ b/examples/basic/vector_ollama.py @@ -0,0 +1,95 @@ +""" +Use CrateDB Vector Search with embeddings computed by a local Ollama server. + +Ollama runs the embedding model in its own process and speaks HTTP, so this +program needs no account, no API key, and no machine learning stack of its +own. `vector_ollama.py` and `vector_openai.py` are otherwise the same program. + +- https://ollama.com/library/nomic-embed-text +- https://python.langchain.com/docs/integrations/text_embedding/ollama/ + +As input data, the example uses the canonical `state_of_the_union.txt`. + +Synopsis:: + + # Install prerequisites. + pip install --upgrade langchain-cratedb langchain-ollama langchain-text-splitters + + # Start database. + docker run --rm -it --publish=4200:4200 crate/crate:nightly + + # Serve the embedding model. `nomic-embed-text` produces 768 dimensions, + # well within the 2048 a CrateDB FLOAT_VECTOR column accepts. + ollama serve + ollama pull nomic-embed-text + + # Optionally set environment variables to configure the Ollama and CrateDB + # endpoints. + export OLLAMA_BASE_URL="http://localhost:11434" + export CRATEDB_SQLALCHEMY_URL="crate://crate@localhost/?schema=doc" + + # Run program. + python examples/basic/vector_ollama.py +""" # noqa: E501 +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "langchain-cratedb", +# "langchain-ollama", +# "langchain-text-splitters", +# ] +# /// + +import os +import typing as t + +import requests +from langchain_core.documents import Document +from langchain_ollama import OllamaEmbeddings +from langchain_text_splitters import RecursiveCharacterTextSplitter + +from langchain_cratedb import CrateDBVectorStore + +CRATEDB_SQLALCHEMY_URL = os.environ.get( + "CRATEDB_SQLALCHEMY_URL", "crate://crate@localhost/?schema=testdrive" +) +OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434") +EMBEDDING_MODEL = "nomic-embed-text" + + +def get_documents() -> t.List[Document]: + """ + Acquire data, return as LangChain documents. + """ + + # Define text splitter. + text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0) + + # Load a document, and split it into chunks. + url = "https://github.com/langchain-ai/langchain/raw/v0.0.325/docs/docs/modules/state_of_the_union.txt" + text = requests.get(url, timeout=10).text + return text_splitter.create_documents([text]) + + +def main() -> None: + # Set up the embedding model. + embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url=OLLAMA_BASE_URL) + + # Acquire documents. + documents = get_documents() + + # Embed each chunk, and load them into the vector store. + vector_store = CrateDBVectorStore.from_documents( + documents=documents, + embedding=embeddings, + connection=CRATEDB_SQLALCHEMY_URL, + ) + + # Invoke a query, and display the first result. + query = "What did the president say about Ketanji Brown Jackson" + docs = vector_store.similarity_search(query) + print(docs[0].page_content) + + +if __name__ == "__main__": + main() diff --git a/examples/basic/vector_search.py b/examples/basic/vector_openai.py similarity index 94% rename from examples/basic/vector_search.py rename to examples/basic/vector_openai.py index 4b5e003..edd6090 100644 --- a/examples/basic/vector_search.py +++ b/examples/basic/vector_openai.py @@ -1,6 +1,8 @@ """ Use CrateDB Vector Search with OpenAI embeddings. +For the same program without an account or an API key, see `vector_ollama.py`. + As input data, the example uses the canonical `state_of_the_union.txt`. Synopsis:: @@ -17,7 +19,7 @@ export CRATEDB_SQLALCHEMY_URL="crate://crate@localhost/?schema=doc" # Run program. - python examples/basic/vector_search.py + python examples/basic/vector_openai.py """ # noqa: E501 # /// script # requires-python = ">=3.9" diff --git a/pyproject.toml b/pyproject.toml index 40cafdf..2f1ba7c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -147,6 +147,9 @@ select = [ omit = [ "langchain_cratedb/retrievers.py", "tests/*", + # Runs against a local Ollama server, which CI does not provide, so the + # suite reaches its import and skips the rest. + "examples/basic/vector_ollama.py", ] [tool.pytest.ini_options] @@ -180,6 +183,7 @@ optional = true [tool.poetry.group.dev.dependencies] [tool.poetry.group.test.dependencies] +langchain-ollama = "<2" langchain-openai = "<1.3" langchain-tests = "==1.1.7" notebook = "<7.6" diff --git a/tests/test_examples.py b/tests/test_examples.py index 404934e..480fa1d 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -35,7 +35,8 @@ def test_file(run_file: t.Callable, file: Path) -> None: """ Execute Python code, one test case per .py file. - Skip test cases that trip when no OpenAI API key is configured. + Skip test cases that need a service this machine does not provide: an + OpenAI API key, or a reachable Ollama server. """ if file.name in SKIP_FILES: raise pytest.skip(f"FIXME: Skipping file: {file.name}") @@ -47,3 +48,7 @@ def test_file(run_file: t.Callable, file: Path) -> None: raise pytest.skip( "Skipping test because `OPENAI_API_KEY` is not defined" ) from ex + except ConnectionError as ex: + if "Failed to connect to Ollama" not in str(ex): + raise + raise pytest.skip("Skipping test because Ollama is not reachable") from ex