Support ollama model

Signed-off-by: Hemslo Wang <[email protected]>

.env.example

@@ -1,2 +1,11 @@
-OPENAI_API_KEY=
 AUTH_TOKEN=
+CHAT_PROVIDER=
+EMBEDDING_DIM=
+EMBEDDING_PROVIDER=
+OLLAMA_CHAT_MODEL=
+OLLAMA_EMBEDDING_MODEL=
+OLLAMA_URL=
+OPENAI_API_KEY=
+OPENAI_CHAT_MODEL=
+OPENAI_EMBEDDING_MODEL=
+REDIS_URL=

app/chains/chat.py

@@ -7,20 +7,23 @@
 from langchain_core.runnables import (
     RunnableParallel,
 )
-from langchain_openai import ChatOpenAI
 from langchain.prompts.prompt import PromptTemplate
 from langserve import CustomUserType
 from pydantic import Field
 
-from app import config
+from app.dependencies.llm import get_llm
 from app.dependencies.redis import get_redis
 
-llm = ChatOpenAI(
-    model=config.OPENAI_CHAT_MODEL,
-    temperature=0,
-)
+llm = get_llm()
 
-retriever = get_redis().as_retriever(search_type="mmr")
+retriever = get_redis().as_retriever(
+    search_type="mmr",
+    search_kwargs={
+        "fetch_k": 20,
+        "k": 3,
+        "lambda_mult": 0.5,
+    },
+)
 
 
 REPHRASE_TEMPLATE = """\

app/config.py

@@ -5,19 +5,20 @@
 
 load_dotenv()
 
+AUTH_TOKEN = os.environ.get("AUTH_TOKEN")
+if AUTH_TOKEN is None:
+    raise ValueError("AUTH_TOKEN is not set in the environment variables")
+
+CHAT_PROVIDER = os.environ.get("CHAT_PROVIDER", "openai")
+EMBEDDING_DIM = int(os.environ.get("EMBEDDING_DIM", 1536))
+EMBEDDING_PROVIDER = os.environ.get("EMBEDDING_PROVIDER", "openai")
+INDEX_NAME = "document"
+INDEX_SCHEMA_PATH = Path(os.path.dirname(__file__)) / "schema.yaml"
+OLLAMA_CHAT_MODEL = os.environ.get("OLLAMA_CHAT_MODEL", "llama2")
+OLLAMA_EMBEDDING_MODEL = os.environ.get("OLLAMA_EMBEDDING_MODEL", "llama2")
+OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
+OPENAI_CHAT_MODEL = os.environ.get("OPENAI_CHAT_MODEL", "gpt-3.5-turbo-1106")
 OPENAI_EMBEDDING_MODEL = os.environ.get(
     "OPENAI_EMBEDDING_MODEL", "text-embedding-3-small"
 )
-
-OPENAI_CHAT_MODEL = os.environ.get("OPENAI_CHAT_MODEL", "gpt-3.5-turbo-1106")
-
-INDEX_SCHEMA_PATH = Path(os.path.dirname(__file__)) / "schema.yaml"
-
 REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/")
-
-INDEX_NAME = "document"
-
-AUTH_TOKEN = os.environ.get("AUTH_TOKEN")
-
-if AUTH_TOKEN is None:
-    raise ValueError("AUTH_TOKEN is not set in the environment variables")

app/dependencies/embeddings.py

@@ -0,0 +1,29 @@
+from langchain_core.embeddings import Embeddings
+
+from app import config
+
+
+def _get_embeddings() -> Embeddings:
+    match config.EMBEDDING_PROVIDER:
+        case "openai":
+            from langchain_openai import OpenAIEmbeddings
+
+            return OpenAIEmbeddings(
+                model=config.OPENAI_EMBEDDING_MODEL,
+            )
+        case "ollama":
+            from langchain_community.embeddings import OllamaEmbeddings
+
+            return OllamaEmbeddings(
+                model=config.OLLAMA_EMBEDDING_MODEL,
+                base_url=config.OLLAMA_URL,
+            )
+        case _:
+            raise ValueError(f"Unknown embedding provider: {config.EMBEDDING_PROVIDER}")
+
+
+embeddings = _get_embeddings()
+
+
+def get_embeddings() -> Embeddings:
+    return embeddings

app/dependencies/llm.py

@@ -0,0 +1,30 @@
+from langchain_core.language_models import BaseChatModel
+
+from app import config
+
+
+def _get_llm() -> BaseChatModel:
+    match config.CHAT_PROVIDER:
+        case "openai":
+            from langchain_openai import ChatOpenAI
+
+            return ChatOpenAI(
+                model=config.OPENAI_CHAT_MODEL,
+                temperature=0,
+            )
+        case "ollama":
+            from langchain_community.chat_models import ChatOllama
+
+            return ChatOllama(
+                model=config.OLLAMA_CHAT_MODEL,
+                base_url=config.OLLAMA_URL,
+            )
+        case _:
+            raise ValueError(f"Unknown chat provider: {config.CHAT_PROVIDER}")
+
+
+llm = _get_llm()
+
+
+def get_llm() -> BaseChatModel:
+    return llm

app/dependencies/redis.py

@@ -2,24 +2,19 @@
 from typing import Annotated
 
 from fastapi import Depends
-from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_community.vectorstores.redis import Redis
 
 from app import config
-
-
-embeddings = OpenAIEmbeddings(
-    model=config.OPENAI_EMBEDDING_MODEL,
-)
+from app.dependencies.embeddings import get_embeddings
 
 rds = Redis(
     redis_url=os.getenv("REDIS_URL", "redis://localhost:6379/"),
     index_name=config.INDEX_NAME,
-    embedding=embeddings,
+    embedding=get_embeddings(),
     index_schema=config.INDEX_SCHEMA_PATH,
 )
 
-rds._create_index_if_not_exist()
+rds._create_index_if_not_exist(config.EMBEDDING_DIM)
 
 
 def get_redis() -> Redis:

app/schema.yaml

@@ -29,12 +29,8 @@ text:
   weight: 1.0
   withsuffixtrie: false
 vector:
-- algorithm: HNSW
+- algorithm: FLAT
   datatype: FLOAT32
   dims: 1536
   distance_metric: COSINE
-  ef_construction: 200
-  ef_runtime: 10
-  epsilon: 0.01
-  m: 16
   name: content_vector

docker-compose.redis.yaml

@@ -0,0 +1,12 @@
+services:
+  redis:
+    image: redis/redis-stack:6.2.6-v11
+    ports:
+      - "6379:6379"
+      - "8001:8001"
+    environment:
+      REDIS_ARGS: --save 60 1000 --appendonly yes
+    volumes:
+      - redis-data:/data
+volumes:
+  redis-data:

docker-compose.yaml

@@ -1,4 +1,5 @@
-version: '3.8'
+include:
+- docker-compose.redis.yaml
 services:
   app:
     build: .
@@ -8,15 +9,5 @@ services:
       REDIS_URL: redis://redis:6379/
     ports:
       - "8000:8000"
-  redis:
-    image: redis/redis-stack:6.2.6-v11
-    ports:
-      - "6379:6379"
-      - "8001:8001"
-    environment:
-      REDIS_ARGS: --save 60 1000 --appendonly yes
-    volumes:
-      - redis-data:/data
-
-volumes:
-    redis-data:
+    depends_on:
+    - redis