Skip to content

Commit 7632736

Browse files
pamelafoxpamelafox
committed
Add embedding with ollama
1 parent d1e990c commit 7632736

File tree

10 files changed

+234033
-942
lines changed

10 files changed

+234033
-942
lines changed

‎.env.sample‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,17 @@ AZURE_OPENAI_CHAT_MODEL=gpt-35-turbo
1818
AZURE_OPENAI_EMBED_DEPLOYMENT=text-embedding-ada-002
1919
AZURE_OPENAI_EMBED_MODEL=text-embedding-ada-002
2020
AZURE_OPENAI_EMBED_MODEL_DIMENSIONS=1536
21+
AZURE_OPENAI_EMBEDDING_COLUMN=embedding_ada002
2122
# Only needed when using key-based Azure authentication:
2223
AZURE_OPENAI_KEY=
2324
# Needed for OpenAI.com:
2425
OPENAICOM_KEY=YOUR-OPENAI-API-KEY
2526
OPENAICOM_CHAT_MODEL=gpt-3.5-turbo
2627
OPENAICOM_EMBED_MODEL=text-embedding-ada-002
2728
OPENAICOM_EMBED_MODEL_DIMENSIONS=1536
29+
OPENAICOM_EMBEDDING_COLUMN=embedding_ada002
2830
# Needed for Ollama:
2931
OLLAMA_ENDPOINT=http://host.docker.internal:11434/v1
30-
OLLAMA_CHAT_MODEL=phi3:3.8b
32+
OLLAMA_CHAT_MODEL=llama3.1
33+
OLLAMA_EMBED_MODEL=nomic-embed-text
34+
OLLAMA_EMBEDDING_COLUMN=embedding_nomic

‎src/backend/fastapi_app/dependencies.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,10 @@ class FastAPIAppContext(BaseModel):
2828

2929
openai_chat_model: str
3030
openai_embed_model: str
31-
openai_embed_dimensions: int
31+
openai_embed_dimensions: int | None
3232
openai_chat_deployment: str | None
3333
openai_embed_deployment: str | None
34+
embedding_column: str
3435

3536

3637
async def common_parameters():
@@ -43,16 +44,24 @@ async def common_parameters():
4344
openai_embed_deployment = os.getenv("AZURE_OPENAI_EMBED_DEPLOYMENT", "text-embedding-ada-002")
4445
openai_embed_model = os.getenv("AZURE_OPENAI_EMBED_MODEL", "text-embedding-ada-002")
4546
openai_embed_dimensions = int(os.getenv("AZURE_OPENAI_EMBED_DIMENSIONS", 1536))
47+
embedding_column = os.getenv("AZURE_OPENAI_EMBEDDING_COLUMN", "embedding_ada002")
48+
elif OPENAI_EMBED_HOST == "ollama":
49+
openai_embed_deployment = None
50+
openai_embed_model = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
51+
openai_embed_dimensions = None
52+
embedding_column = os.getenv("OLLAMA_EMBEDDING_COLUMN", "embedding_nomic")
4653
else:
47-
openai_embed_deployment = "text-embedding-ada-002"
54+
openai_embed_deployment = None
4855
openai_embed_model = os.getenv("OPENAICOM_EMBED_MODEL", "text-embedding-ada-002")
4956
openai_embed_dimensions = int(os.getenv("OPENAICOM_EMBED_DIMENSIONS", 1536))
57+
embedding_column = os.getenv("OPENAICOM_EMBEDDING_COLUMN", "embedding_ada002")
5058
if OPENAI_CHAT_HOST == "azure":
5159
openai_chat_deployment = os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT", "gpt-35-turbo")
5260
openai_chat_model = os.getenv("AZURE_OPENAI_CHAT_MODEL", "gpt-35-turbo")
5361
elif OPENAI_CHAT_HOST == "ollama":
5462
openai_chat_deployment = None
5563
openai_chat_model = os.getenv("OLLAMA_CHAT_MODEL", "phi3:3.8b")
64+
openai_embed_model = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
5665
else:
5766
openai_chat_deployment = None
5867
openai_chat_model = os.getenv("OPENAICOM_CHAT_MODEL", "gpt-3.5-turbo")
@@ -62,6 +71,7 @@ async def common_parameters():
6271
openai_embed_dimensions=openai_embed_dimensions,
6372
openai_chat_deployment=openai_chat_deployment,
6473
openai_embed_deployment=openai_embed_deployment,
74+
embedding_column=embedding_column,
6575
)
6676

6777

‎src/backend/fastapi_app/embeddings.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@ async def compute_text_embedding(
2121
class ExtraArgs(TypedDict, total=False):
2222
dimensions: int
2323

24-
dimensions_args: ExtraArgs = {"dimensions": embedding_dimensions} if SUPPORTED_DIMENSIONS_MODEL[embed_model] else {}
24+
dimensions_args: ExtraArgs = (
25+
{"dimensions": embedding_dimensions} if SUPPORTED_DIMENSIONS_MODEL.get(embed_model) else {}
26+
)
2527

2628
embedding = await openai_client.embeddings.create(
2729
# Azure OpenAI takes the deployment name as the model name

‎src/backend/fastapi_app/openai_clients.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,13 @@ async def create_openai_embed_client(
7676
azure_deployment=azure_deployment,
7777
azure_ad_token_provider=token_provider,
7878
)
79-
79+
elif OPENAI_EMBED_HOST == "ollama":
80+
logger.info("Authenticating to OpenAI using Ollama...")
81+
openai_embed_client = openai.AsyncOpenAI(
82+
base_url=os.getenv("OLLAMA_ENDPOINT"),
83+
api_key="nokeyneeded",
84+
)
8085
else:
86+
logger.info("Authenticating to OpenAI using OpenAI.com API key...")
8187
openai_embed_client = openai.AsyncOpenAI(api_key=os.getenv("OPENAICOM_KEY"))
8288
return openai_embed_client

‎src/backend/fastapi_app/postgres_models.py‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,17 @@ class Item(Base):
2020
name: Mapped[str] = mapped_column()
2121
description: Mapped[str] = mapped_column()
2222
price: Mapped[float] = mapped_column()
23-
embedding: Mapped[Vector] = mapped_column(Vector(1536)) # ada-002
23+
embedding_ada002: Mapped[Vector] = mapped_column(Vector(1536)) # ada-002
24+
embedding_nomic: Mapped[Vector] | None = mapped_column(Vector(768), nullable=True) # nomic-embed-text
2425

2526
def to_dict(self, include_embedding: bool = False):
2627
model_dict = asdict(self)
2728
if include_embedding:
28-
model_dict["embedding"] = model_dict["embedding"].tolist()
29+
model_dict["embedding_ada002"] = model_dict.get("embedding_ada002", [])
30+
model_dict["embedding_nomic"] = model_dict.get("embedding_nomic", [])
2931
else:
30-
del model_dict["embedding"]
32+
del model_dict["embedding_ada002"]
33+
del model_dict["embedding_nomic"]
3134
return model_dict
3235

3336
def to_str_for_rag(self):
@@ -38,10 +41,18 @@ def to_str_for_embedding(self):
3841

3942

4043
# Define HNSW index to support vector similarity search through the vector_cosine_ops access method (cosine distance).
41-
index = Index(
42-
"hnsw_index_for_innerproduct_item_embedding",
43-
Item.embedding,
44+
index_ada002 = Index(
45+
"hnsw_index_for_innerproduct_item_embedding_ada002",
46+
Item.embedding_ada002,
4447
postgresql_using="hnsw",
4548
postgresql_with={"m": 16, "ef_construction": 64},
46-
postgresql_ops={"embedding": "vector_ip_ops"},
49+
postgresql_ops={"embedding_ada002": "vector_ip_ops"},
50+
)
51+
52+
index_nomic = Index(
53+
"hnsw_index_for_innerproduct_item_embedding_nomic",
54+
Item.embedding_nomic,
55+
postgresql_using="hnsw",
56+
postgresql_with={"m": 16, "ef_construction": 64},
57+
postgresql_ops={"embedding_nomic": "vector_ip_ops"},
4758
)

‎src/backend/fastapi_app/postgres_searcher.py‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,14 @@ def __init__(
1515
embed_deployment: str | None, # Not needed for non-Azure OpenAI or for retrieval_mode="text"
1616
embed_model: str,
1717
embed_dimensions: int,
18+
embedding_column: str,
1819
):
1920
self.db_session = db_session
2021
self.openai_embed_client = openai_embed_client
2122
self.embed_model = embed_model
2223
self.embed_deployment = embed_deployment
2324
self.embed_dimensions = embed_dimensions
25+
self.embedding_column = embedding_column
2426

2527
def build_filter_clause(self, filters) -> tuple[str, str]:
2628
if filters is None:
@@ -36,19 +38,15 @@ def build_filter_clause(self, filters) -> tuple[str, str]:
3638
return "", ""
3739

3840
async def search(
39-
self,
40-
query_text: str | None,
41-
query_vector: list[float] | list,
42-
top: int = 5,
43-
filters: list[dict] | None = None,
41+
self, query_text: str | None, query_vector: list[float] | list, top: int = 5, filters: list[dict] | None = None
4442
):
4543
filter_clause_where, filter_clause_and = self.build_filter_clause(filters)
4644

4745
vector_query = f"""
48-
SELECT id, RANK () OVER (ORDER BY embedding <=> :embedding) AS rank
46+
SELECT id, RANK () OVER (ORDER BY {self.embedding_column} <=> :embedding) AS rank
4947
FROM items
5048
{filter_clause_where}
51-
ORDER BY embedding <=> :embedding
49+
ORDER BY {self.embedding_column} <=> :embedding
5250
LIMIT 20
5351
"""
5452

‎src/backend/fastapi_app/routes/api_routes.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ async def chat_handler(
9999
embed_deployment=context.openai_embed_deployment,
100100
embed_model=context.openai_embed_model,
101101
embed_dimensions=context.openai_embed_dimensions,
102+
embedding_column=context.embedding_column,
102103
)
103104
rag_flow: SimpleRAGChat | AdvancedRAGChat
104105
if chat_request.context.overrides.use_advanced_flow:
@@ -139,6 +140,7 @@ async def chat_stream_handler(
139140
embed_deployment=context.openai_embed_deployment,
140141
embed_model=context.openai_embed_model,
141142
embed_dimensions=context.openai_embed_dimensions,
143+
embedding_column=context.embedding_column,
142144
)
143145

144146
rag_flow: SimpleRAGChat | AdvancedRAGChat

‎src/backend/fastapi_app/seed_data.json‎

Lines changed: 233915 additions & 908 deletions
Large diffs are not rendered by default.

‎src/backend/fastapi_app/setup_postgres_seeddata.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,17 +36,18 @@ async def seed_data(engine):
3636
with open(os.path.join(current_dir, "seed_data.json")) as f:
3737
catalog_items = json.load(f)
3838
for catalog_item in catalog_items:
39-
db_item = await session.execute(select(Item).filter(Item.id == catalog_item["Id"]))
39+
db_item = await session.execute(select(Item).filter(Item.id == catalog_item["id"]))
4040
if db_item.scalars().first():
4141
continue
4242
item = Item(
43-
id=catalog_item["Id"],
44-
type=catalog_item["Type"],
45-
brand=catalog_item["Brand"],
46-
name=catalog_item["Name"],
47-
description=catalog_item["Description"],
48-
price=catalog_item["Price"],
49-
embedding=catalog_item["Embedding"],
43+
id=catalog_item["id"],
44+
type=catalog_item["type"],
45+
brand=catalog_item["brand"],
46+
name=catalog_item["name"],
47+
description=catalog_item["description"],
48+
price=catalog_item["price"],
49+
embedding_ada002=catalog_item["embedding_ada002"],
50+
embedding_nomic=catalog_item.get("embedding_nomic"),
5051
)
5152
session.add(item)
5253
try:

‎src/backend/fastapi_app/update_embeddings.py‎

Lines changed: 57 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
11
import asyncio
2+
import json
3+
import logging
4+
import os
25

36
from dotenv import load_dotenv
47
from sqlalchemy import select
@@ -10,28 +13,75 @@
1013
from fastapi_app.postgres_engine import create_postgres_engine_from_env
1114
from fastapi_app.postgres_models import Item
1215

16+
logger = logging.getLogger("ragapp")
1317

14-
async def update_embeddings():
18+
19+
async def update_embeddings(in_seed_data=False):
1520
azure_credential = await get_azure_credentials()
1621
engine = await create_postgres_engine_from_env(azure_credential)
1722
openai_embed_client = await create_openai_embed_client(azure_credential)
1823
common_params = await common_parameters()
1924

20-
async with async_sessionmaker(engine, expire_on_commit=False)() as session:
21-
async with session.begin():
22-
items = (await session.scalars(select(Item))).all()
23-
24-
for item in items:
25-
item.embedding = await compute_text_embedding(
25+
embedding_column = ""
26+
OPENAI_EMBED_HOST = os.getenv("OPENAI_EMBED_HOST")
27+
if OPENAI_EMBED_HOST == "azure":
28+
embedding_column = os.getenv("AZURE_OPENAI_EMBEDDING_COLUMN", "embedding_ada002")
29+
elif OPENAI_EMBED_HOST == "ollama":
30+
embedding_column = os.getenv("OLLAMA_EMBEDDING_COLUMN", "embedding_nomic")
31+
else:
32+
embedding_column = os.getenv("OPENAICOM_EMBEDDING_COLUMN", "embedding_ada002")
33+
logger.info(f"Updating embeddings in column: {embedding_column}")
34+
if in_seed_data:
35+
current_dir = os.path.dirname(os.path.realpath(__file__))
36+
items = []
37+
with open(os.path.join(current_dir, "seed_data.json")) as f:
38+
catalog_items = json.load(f)
39+
for catalog_item in catalog_items:
40+
item = Item(
41+
id=catalog_item["id"],
42+
type=catalog_item["type"],
43+
brand=catalog_item["brand"],
44+
name=catalog_item["name"],
45+
description=catalog_item["description"],
46+
price=catalog_item["price"],
47+
embedding_ada002=catalog_item["embedding_ada002"],
48+
embedding_nomic=catalog_item.get("embedding_nomic"),
49+
)
50+
embedding = await compute_text_embedding(
2651
item.to_str_for_embedding(),
2752
openai_client=openai_embed_client,
2853
embed_model=common_params.openai_embed_model,
54+
embed_deployment=common_params.openai_embed_deployment,
2955
embedding_dimensions=common_params.openai_embed_dimensions,
3056
)
57+
setattr(item, embedding_column, embedding)
58+
items.append(item)
59+
# write to the file
60+
with open(os.path.join(current_dir, "seed_data.json"), "w") as f:
61+
json.dump([item.to_dict(include_embedding=True) for item in items], f, indent=4)
62+
return
3163

64+
async with async_sessionmaker(engine, expire_on_commit=False)() as session:
65+
async with session.begin():
66+
items = (await session.scalars(select(Item))).all()
67+
68+
for item in items:
69+
setattr(
70+
item,
71+
embedding_column,
72+
await compute_text_embedding(
73+
item.to_str_for_embedding(),
74+
openai_client=openai_embed_client,
75+
embed_model=common_params.openai_embed_model,
76+
embed_deployment=common_params.openai_embed_deployment,
77+
embedding_dimensions=common_params.openai_embed_dimensions,
78+
),
79+
)
3280
await session.commit()
3381

3482

3583
if __name__ == "__main__":
84+
logging.basicConfig(level=logging.WARNING)
85+
logger.setLevel(logging.INFO)
3686
load_dotenv(override=True)
3787
asyncio.run(update_embeddings())

0 commit comments

Comments
 (0)