Add support for local models via Ollama

Ollama.md

@@ -0,0 +1,14 @@
+# Ollama setup
+1. Download and install [Ollama](https://ollama.com/)
+2. Once Ollama is running on your system, run `ollama pull llama3.1`
+> Currently this is a ~5GB download, it's best to download it before the workshop if you plan on using it
+3. Update the `MODEL_NAME` in your `dot.env` file to `ollama`
+
+You're now ready to begin the workshop! Head back to the [Readme.md](Readme.md)
+
+## Restarting the workshop 
+Mixing use of llama and openai on the same Redis instance can cause unexpected behavior. If you want to switch from one to the other it is recommended to kill and re-create the instance. To do this:
+1. Run `docker ps` and take note of the ID for the running image
+2. `docker stop imageId`
+3. `docker rm imageId`
+4. Start a new instance using the command from earlier, `docker run -d --name redis -p 6379:6379 -p 8001:8001 redis/redis-stack:latest`

Readme.md

@@ -17,6 +17,10 @@ In this workshop, we are going to use [LangGraph](https://langchain-ai.github.io
 - [docker](https://docs.docker.com/get-started/get-docker/)
 - [openai api key](https://platform.openai.com/docs/quickstart)
 
+## (Optional) Ollama
+This workshop is optimized to run targeting OpenAI models. If you prefer to run locally however, you may do so via the experimental Ollama configuration. 
+* [Ollama setup instructions](Ollama.md)
+
 ## (Optional) helpers
 
 - [LangSmith](https://docs.smith.langchain.com/)
@@ -235,7 +239,11 @@ In our scenario we want to be able to retrieve the time-bound information that t
 
 ### Steps:
 - Open [participant_agent/utils/vector_store.py](participant_agent/utils/vector_store.py)
+- Take note of how `embedding_model` is getting instantiated. If using Ollama then switch this for the appropriate embedding using `llama3.1` for the `model` parameter
+> [OpenAI embeddings](https://python.langchain.com/docs/integrations/text_embedding/openai/) \
+[Ollama embeddings](https://python.langchain.com/docs/integrations/text_embedding/ollama/)
 - Where `vector_store=None` update to `vector_store = RedisVectorStore.from_documents(<docs>, <embedding_model>, config=<config>)` with the appropriate variables.
+
 - Open [participant_agent/utils/tools.py](participant_agent/utils/tools.py)
     - Uncomment code for retrieval tool
     - Update the create_retriever_tool to take the correct params. Ex: `create_retriever_tool(vector_store.as_retriever(), "get_directions", "meaningful doc string")`

dot.env

@@ -3,4 +3,5 @@ OPENAI_API_KEY=openai_key
 LANGCHAIN_TRACING_V2=
 LANGCHAIN_ENDPOINT=
 LANGCHAIN_API_KEY=
-LANGCHAIN_PROJECT=
+LANGCHAIN_PROJECT=
+MODEL_NAME=openai

example_agent/ex_graph.py

@@ -11,7 +11,7 @@
 
 # Define the config
 class GraphConfig(TypedDict):
-    model_name: Literal["anthropic", "openai"]
+    model_name: Literal["anthropic", "openai", "ollama"]
 
 
 # Define the function that determines whether to continue or not

example_agent/utils/ex_nodes.py

@@ -1,18 +1,26 @@
+import os
 from functools import lru_cache
 
+from dotenv import load_dotenv
 from langchain_core.messages import HumanMessage
 from langchain_openai import ChatOpenAI
+from langchain_ollama import ChatOllama
 from langgraph.prebuilt import ToolNode
 
 from example_agent.utils.ex_tools import tools
 
 from .ex_state import AgentState, MultipleChoiceResponse
 
+load_dotenv()
+
+ENVIRON_MODEL_NAME = os.environ.get("MODEL_NAME")
 
 @lru_cache(maxsize=4)
 def _get_tool_model(model_name: str):
     if model_name == "openai":
         model = ChatOpenAI(temperature=0, model_name="gpt-4o")
+    elif model_name == "ollama":
+        model = ChatOllama(temperature=0, model="llama3.1", num_ctx=4096)
     else:
         raise ValueError(f"Unsupported model type: {model_name}")
 
@@ -24,6 +32,8 @@ def _get_tool_model(model_name: str):
 def _get_response_model(model_name: str):
     if model_name == "openai":
         model = ChatOpenAI(temperature=0, model_name="gpt-4o")
+    elif model_name == "ollama":
+        model = ChatOllama(temperature=0, model="llama3.1", num_ctx=4096)
     else:
         raise ValueError(f"Unsupported model type: {model_name}")
 
@@ -36,7 +46,7 @@ def multi_choice_structured(state: AgentState, config):
     # We call the model with structured output in order to return the same format to the user every time
     # state['messages'][-2] is the last ToolMessage in the convo, which we convert to a HumanMessage for the model to use
     # We could also pass the entire chat history, but this saves tokens since all we care to structure is the output of the tool
-    model_name = config.get("configurable", {}).get("model_name", "openai")
+    model_name = config.get("configurable", {}).get("model_name", ENVIRON_MODEL_NAME)
 
     response = _get_response_model(model_name).invoke(
         [
@@ -75,7 +85,7 @@ def call_tool_model(state: AgentState, config):
     messages = [{"role": "system", "content": system_prompt}] + state["messages"]
 
     # Get from LangGraph config
-    model_name = config.get("configurable", {}).get("model_name", "openai")
+    model_name = config.get("configurable", {}).get("model_name", ENVIRON_MODEL_NAME)
 
     # Get our model that binds our tools
     model = _get_tool_model(model_name)

example_agent/utils/ex_vector_store.py

@@ -3,6 +3,8 @@
 from dotenv import load_dotenv
 from langchain_core.documents import Document
 from langchain_openai import OpenAIEmbeddings
+from langchain_ollama import OllamaEmbeddings
+from redis import Redis
 from langchain_redis import RedisConfig, RedisVectorStore
 
 load_dotenv()
@@ -11,20 +13,30 @@
 INDEX_NAME = os.environ.get("VECTOR_INDEX_NAME", "oregon_trail")
 
 config = RedisConfig(index_name=INDEX_NAME, redis_url=REDIS_URL)
+redis_client = Redis.from_url(REDIS_URL)
 
 doc = Document(
     page_content="the northern trail, of the blue mountains, was destroyed by a flood and is no longer safe to traverse. It is recommended to take the southern trail although it is longer."
 )
 
+ # TODO: participant can change to whatever desired model
+embedding_model = OpenAIEmbeddings()
+# embedding_model = OllamaEmbeddings(model="llama3.1")
+
+def _clean_existing(prefix):
+    for key in redis_client.scan_iter(f"{prefix}:*"):
+        redis_client.delete(key)
 
 def get_vector_store():
     try:
         config.from_existing = True
-        vector_store = RedisVectorStore(OpenAIEmbeddings(), config=config)
+        vector_store = RedisVectorStore(embedding_model, config=config)
     except:
         print("Init vector store with document")
+        print("Clean any existing data in index")
+        _clean_existing(config.index_name)
         config.from_existing = False
         vector_store = RedisVectorStore.from_documents(
-            [doc], OpenAIEmbeddings(), config=config
+            [doc], embedding_model, config=config
         )
     return vector_store

participant_agent/graph.py

@@ -14,7 +14,7 @@
 
 # The graph config can be updated with LangGraph Studio which can be helpful
 class GraphConfig(TypedDict):
-    model_name: Literal["openai"]  # could add more LLM providers here
+    model_name: Literal["openai", "ollama"]  # could add more LLM providers here
 
 
 # Define the function that determines whether to continue or not

participant_agent/utils/nodes.py

@@ -1,13 +1,18 @@
+import os
 from functools import lru_cache
 
+from dotenv import load_dotenv
 from langchain_core.messages import HumanMessage
 from langchain_openai import ChatOpenAI
+from langchain_ollama import ChatOllama
 from langgraph.prebuilt import ToolNode
 
 from participant_agent.utils.tools import tools
 
 from .state import AgentState, MultipleChoiceResponse
 
+load_dotenv()
+
 
 # need to use this in call_tool_model function
 @lru_cache(maxsize=4)
@@ -17,6 +22,8 @@ def _get_tool_model(model_name: str):
     """
     if model_name == "openai":
         model = ChatOpenAI(temperature=0, model_name="gpt-4o")
+    elif model_name == "ollama":
+        model = ChatOllama(temperature=0, model="llama3.1", num_ctx=4096)
     else:
         raise ValueError(f"Unsupported model type: {model_name}")
 
@@ -32,6 +39,8 @@ def _get_tool_model(model_name: str):
 def _get_response_model(model_name: str):
     if model_name == "openai":
         model = ChatOpenAI(temperature=0, model_name="gpt-4o")
+    elif model_name == "ollama":
+        model = ChatOllama(temperature=0, model="llama3.1", num_ctx=4096)
     else:
         raise ValueError(f"Unsupported model type: {model_name}")
 
@@ -45,7 +54,7 @@ def multi_choice_structured(state: AgentState, config):
     # We call the model with structured output in order to return the same format to the user every time
     # state['messages'][-2] is the last ToolMessage in the convo, which we convert to a HumanMessage for the model to use
     # We could also pass the entire chat history, but this saves tokens since all we care to structure is the output of the tool
-    model_name = config.get("configurable", {}).get("model_name", "openai")
+    model_name = config.get("configurable", {}).get("model_name", os.environ.get("MODEL_NAME"))
 
     response = _get_response_model(model_name).invoke(
         [
@@ -84,8 +93,7 @@ def call_tool_model(state: AgentState, config):
     messages = [{"role": "system", "content": system_prompt}] + state["messages"]
 
     # Get from LangGraph config
-    model_name = config.get("configurable", {}).get("model_name", "openai")
-
+    model_name = config.get("configurable", {}).get("model_name", os.environ.get("MODEL_NAME"))
     # Get our model that binds our tools
     model = _get_tool_model(model_name)
 

participant_agent/utils/vector_store.py

@@ -3,6 +3,7 @@
 from dotenv import load_dotenv
 from langchain_core.documents import Document
 from langchain_openai import OpenAIEmbeddings
+from langchain_ollama import OllamaEmbeddings
 from langchain_redis import RedisConfig, RedisVectorStore
 
 load_dotenv()
@@ -11,18 +12,27 @@
 INDEX_NAME = os.environ.get("VECTOR_INDEX_NAME", "oregon_trail")
 
 config = RedisConfig(index_name=INDEX_NAME, redis_url=REDIS_URL)
+redis_client = Redis.from_url(REDIS_URL)
 
 doc = Document(
     page_content="the northern trail, of the blue mountains, was destroyed by a flood and is no longer safe to traverse. It is recommended to take the southern trail although it is longer."
 )
 
+# TODO: participant can change to whatever desired model
+embedding_model = OpenAIEmbeddings() 
+
+def _clean_existing(prefix):
+    for key in redis_client.scan_iter(f"{prefix}:*"):
+        redis_client.delete(key)
 
 def get_vector_store():
     try:
         config.from_existing = True
-        vector_store = RedisVectorStore(OpenAIEmbeddings(), config=config)
+        vector_store = RedisVectorStore(embedding_model, config=config)
     except:
         print("Init vector store with document")
+        print("Clean any existing data in index")
+        _clean_existing(config.index_name)
         config.from_existing = False
 
         # TODO: define vector store

requirements.txt

@@ -1,6 +1,7 @@
 langgraph==0.2.56
 langchain==0.3.13
 langchain-openai==0.2.3
+langchain-ollama==0.2.3
 langchain-redis==0.1.1
 pydantic==2.9.2
 python-dotenv==1.0.1

test_example_oregon_trail.py

@@ -35,7 +35,7 @@ def test_1_wagon_leader(app):
 
     res = graph.invoke({"messages": scenario["question"]})
 
-    assert res["messages"][-1].content == scenario["answer"]
+    assert scenario["answer"] in res["messages"][-1].content
 
     print(f"\n response: {scenario['answer']}")
 

test_participant_oregon_trail.py

@@ -33,7 +33,7 @@ def test_1_wagon_leader(app):
 
     res = graph.invoke({"messages": scenario["question"]})
 
-    assert res["messages"][-1].content == scenario["answer"]
+    assert scenario["answer"] in res["messages"][-1].content
 
     print(f"\n response: {scenario['answer']}")
 

test_setup.py

@@ -2,11 +2,18 @@
 
 from dotenv import load_dotenv
 from langchain_openai import ChatOpenAI
+from langchain_ollama import ChatOllama
 from redis import Redis
 
 load_dotenv()
 
-llm = ChatOpenAI(model="gpt-4o")
+if os.environ.get("MODEL_NAME") == "openai":
+    llm = ChatOpenAI(model="gpt-4o")
+elif os.environ.get("MODEL_NAME") == "ollama":
+    llm = ChatOllama(model="llama3.1")
+else:
+    raise Exception("Setup failed, MODEL_NAME not defined in .env")
+
 client = Redis.from_url(os.environ.get("REDIS_URL"))