feat: Add support for Ollama

.env.example

@@ -1,4 +1,7 @@
 OPENAI_API_KEY="sk-***"
+# OPENAI_API_MODEL="" # Uncomment to change the default OPENAI_API model ("gpt-4.1") and specify a new one
+# OLLAMA_BASE_URL="http://localhost:11434" # Comment OPENAI_API_KEY and uncomment OLLAMA_BASE_URL to use Ollama
+# OLLAMA_MODEL="" # Uncomment to change the default OLLAMA model ("llama2") and specify a new one
 LLAMACLOUD_API_KEY="llx-***"
 ELEVENLABS_API_KEY="sk_***"
 pgql_db="postgres"

README.md

@@ -65,6 +65,15 @@ Next, open the `.env` file and add your API keys:
 - `ELEVENLABS_API_KEY`: find it [on ElevenLabs Settings](https://elevenlabs.io/app/settings/api-keys)
 - `LLAMACLOUD_API_KEY`: find it [on LlamaCloud Dashboard](https://cloud.llamaindex.ai?utm_source=demo&utm_medium=notebookLM)
 
+To use Ollama instead of the OpenAI API, add your base URL:
+
+- `OLLAMA_BASE_URL`: by default is "http://localhost:11434"
+
+Optionally, you can also add the model:
+
+- `OPENAI_API_MODEL`: by default is "gpt-4.1"
+- `OLLAMA_MODEL`: by default is "llama2"
+
 **4. Activate the Virtual Environment**
 
 (on mac/unix)

src/notebookllama/audio.py

@@ -13,6 +13,7 @@
 from pydantic import BaseModel, ConfigDict, model_validator, Field
 from llama_index.core.llms import ChatMessage
 from llama_index.llms.openai import OpenAIResponses
+from llama_index.llms.ollama import Ollama
 
 logger = logging.getLogger(__name__)
 
@@ -362,12 +363,24 @@ async def create_conversation(
 
 PODCAST_GEN: Optional[PodcastGenerator]
 
-if os.getenv("ELEVENLABS_API_KEY", None) and os.getenv("OPENAI_API_KEY", None):
-    SLLM = OpenAIResponses(
-        model="gpt-4.1", api_key=os.getenv("OPENAI_API_KEY")
-    ).as_structured_llm(MultiTurnConversation)
+SLLM = None
+EL_CLIENT = None
+PODCAST_GEN = None
+if os.getenv("ELEVENLABS_API_KEY", None):
     EL_CLIENT = AsyncElevenLabs(api_key=os.getenv("ELEVENLABS_API_KEY"))
-    PODCAST_GEN = PodcastGenerator(llm=SLLM, client=EL_CLIENT)
+    if os.getenv("OPENAI_API_KEY", None):
+        SLLM = OpenAIResponses(
+            model=os.getenv("OPENAI_API_MODEL", "gpt-4.1"),
+            api_key=os.getenv("OPENAI_API_KEY")
+        ).as_structured_llm(MultiTurnConversation)
+    elif os.getenv("OLLAMA_BASE_URL", None):
+        SLLM = Ollama(
+            model=os.getenv("OLLAMA_MODEL", "llama2"),
+            base_url=os.getenv("OLLAMA_BASE_URL")
+        ).as_structured_llm(MultiTurnConversation)
+    if SLLM and EL_CLIENT:
+        PODCAST_GEN = PodcastGenerator(llm=SLLM, client=EL_CLIENT)
+    else:
+        logger.warning("Missing LLM or ElevenLabs client - PODCAST_GEN not initialized")
 else:
-    logger.warning("Missing API keys - PODCAST_GEN not initialized")
-    PODCAST_GEN = None
+    logger.warning("Missing ElevenLabs API key - PODCAST_GEN not initialized")

src/notebookllama/mindmap.py

@@ -1,3 +1,4 @@
+import logging
 import uuid
 import os
 import warnings
@@ -9,7 +10,9 @@
 from pyvis.network import Network
 from llama_index.core.llms import ChatMessage
 from llama_index.llms.openai import OpenAIResponses
+from llama_index.llms.ollama import Ollama
 
+logger = logging.getLogger(__name__)
 
 class Node(BaseModel):
     id: str
@@ -68,9 +71,22 @@ class MindMapCreationFailedWarning(Warning):
     """A warning returned if the mind map creation failed"""
 
 
+LLM = None
+LLM_STRUCT = None
 if os.getenv("OPENAI_API_KEY", None):
-    LLM = OpenAIResponses(model="gpt-4.1", api_key=os.getenv("OPENAI_API_KEY"))
+    LLM = OpenAIResponses(
+        model=os.getenv("OPENAI_API_MODEL", "gpt-4.1"),
+        api_key=os.getenv("OPENAI_API_KEY")
+    )
+elif os.getenv("OLLAMA_BASE_URL", None):
+    LLM = Ollama(
+        model=os.getenv("OLLAMA_MODEL", "llama2"),
+        base_url=os.getenv("OLLAMA_BASE_URL")
+    )
+if LLM:
     LLM_STRUCT = LLM.as_structured_llm(MindMap)
+else:
+    logger.warning("Missing LLM - MindMap generator not initialized")
 
 
 async def get_mind_map(summary: str, highlights: List[str]) -> Union[str, None]:

src/notebookllama/querying.py

@@ -1,30 +1,50 @@
 from dotenv import load_dotenv
+import logging
 import os
 
 from llama_index.core.query_engine import CitationQueryEngine
 from llama_index.core.base.response.schema import Response
 from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
 from llama_index.llms.openai import OpenAIResponses
+from llama_index.llms.ollama import Ollama
 from typing import Union, cast
 
+logger = logging.getLogger(__name__)
 load_dotenv()
 
+LLM = None
+RETR = None
+QE = None
+
 if (
     os.getenv("LLAMACLOUD_API_KEY", None)
     and os.getenv("LLAMACLOUD_PIPELINE_ID", None)
-    and os.getenv("OPENAI_API_KEY", None)
 ):
-    LLM = OpenAIResponses(model="gpt-4.1", api_key=os.getenv("OPENAI_API_KEY"))
-    PIPELINE_ID = os.getenv("LLAMACLOUD_PIPELINE_ID")
     RETR = LlamaCloudIndex(
-        api_key=os.getenv("LLAMACLOUD_API_KEY"), pipeline_id=PIPELINE_ID
+        api_key=os.getenv("LLAMACLOUD_API_KEY"),
+        pipeline_id=os.getenv("LLAMACLOUD_PIPELINE_ID", None)
     ).as_retriever()
+
+if os.getenv("OPENAI_API_KEY", None):
+    LLM = OpenAIResponses(
+        model=os.getenv("OPENAI_API_MODEL", "gpt-4.1"),
+        api_key=os.getenv("OPENAI_API_KEY")
+    )
+elif os.getenv("OLLAMA_BASE_URL", None):
+    LLM = Ollama(
+        model=os.getenv("OLLAMA_MODEL", "llama2"),
+        base_url=os.getenv("OLLAMA_BASE_URL")
+    )
+
+if RETR and LLM:
     QE = CitationQueryEngine(
         retriever=RETR,
         llm=LLM,
         citation_chunk_size=256,
         citation_chunk_overlap=50,
     )
+else:
+    logger.warning("Missing LLM or retriever - Query Engine not initialized")
 
 
 async def query_index(question: str) -> Union[str, None]:

src/notebookllama/verifying.py

@@ -1,13 +1,16 @@
 from dotenv import load_dotenv
+import logging
 import json
 import os
 
 from pydantic import BaseModel, Field, model_validator
 from llama_index.core.llms import ChatMessage
 from llama_index.llms.openai import OpenAIResponses
+from llama_index.llms.ollama import Ollama
 from typing import List, Tuple, Optional
 from typing_extensions import Self
 
+logger = logging.getLogger(__name__)
 load_dotenv()
 
 
@@ -29,9 +32,22 @@ def validate_claim_ver(self) -> Self:
         return self
 
 
+LLM = None
+LLM_VERIFIER = None
 if os.getenv("OPENAI_API_KEY", None):
-    LLM = OpenAIResponses(model="gpt-4.1", api_key=os.getenv("OPENAI_API_KEY"))
+    LLM = OpenAIResponses(
+        model=os.getenv("OPENAI_API_MODEL", "gpt-4.1"),
+        api_key=os.getenv("OPENAI_API_KEY")
+    )
+elif os.getenv("OLLAMA_BASE_URL", None):
+    LLM = Ollama(
+        model=os.getenv("OLLAMA_MODEL", "llama2"),
+        base_url=os.getenv("OLLAMA_BASE_URL")
+    )
+if LLM:
     LLM_VERIFIER = LLM.as_structured_llm(ClaimVerification)
+else:
+    logger.warning("Missing LLM - Verifier not initialized")
 
 
 def verify_claim(

tests/test_utils.py

@@ -22,7 +22,10 @@
     os.getenv("LLAMACLOUD_API_KEY", None)
     and os.getenv("EXTRACT_AGENT_ID", None)
     and os.getenv("LLAMACLOUD_PIPELINE_ID", None)
-    and os.getenv("OPENAI_API_KEY", None)
+    and (
+        os.getenv("OPENAI_API_KEY", None)
+        or os.getenv("OLLAMA_BASE_URL", None)
+    )
 )