feat: implement get chat completions APIs #2200

ehhuang · 2025-05-19T05:20:33Z

What does this PR do?

Provide sqlite implementation of the APIs introduced in feat: introduce APIs for retrieving chat completion requests #2145.
Introduced a SqlStore API: llama_stack/providers/utils/sqlstore/api.py and the first Sqlite implementation
Pagination support will be added in a future PR.

Test Plan

INFERENCE_MODEL="llama3.2:3b-instruct-fp16" llama stack build --template ollama --image-type conda --run

LLAMA_STACK_CONFIG=http://localhost:5001 INFERENCE_MODEL="llama3.2:3b-instruct-fp16" python -m pytest -v tests/integration/inference/test_openai_completion.py --text-model "llama3.2:3b-instruct-fp16" -k 'inference_store and openai'

llama_stack/providers/utils/inference/inference_store.py

leseb · 2025-05-19T12:46:44Z

llama_stack/distribution/routers/__init__.py

@@ -76,6 +80,9 @@ async def get_auto_router_impl(api: Api, routing_table: RoutingTable, deps: dict
        if dep_api in deps:
            api_to_dep_impl[dep_name] = deps[dep_api]

+    if api.value == "inference" and run_config.inference_store:


Suggested change

if api.value == "inference" and run_config.inference_store:

if api.value == Api.inference and run_config.inference_store:

llama_stack/distribution/routers/routers.py

leseb · 2025-05-19T12:50:59Z

llama_stack/distribution/routers/routers.py

+        if self.store:
+            return await self.store.list_chat_completions(after, limit, model, order)
+        else:
+            raise NotImplementedError("List chat completions not implemented")


Is it an error, or just that the store is not configured? The config is optional, right?

updated the message.

leseb · 2025-05-19T12:51:08Z

llama_stack/distribution/routers/routers.py

+        if self.store:
+            return await self.store.get_chat_completion(completion_id)
+        else:
+            raise NotImplementedError("Get chat completion not implemented")


leseb · 2025-05-19T12:51:43Z

llama_stack/providers/utils/inference/inference_store.py

+
+        impl = SqliteInferenceStore(config.db_path)
+    else:
+        raise ValueError(f"Unknown inference store type {config.type}")


Print the supported types as well?

leseb · 2025-05-19T12:53:37Z

llama_stack/providers/utils/inference/stores/sqlite.py

+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(self.conn_string), exist_ok=True)
+
+        async with aiosqlite.connect(self.conn_string) as conn:


It seems like most of the code in this file already exists and is duplicated from other vector db implementations. Is there any chance to reuse some of those instead of adding another instance?

At a quick glance, only the table name would vary.

The sql queries and the surrounding code is the bulk of the logic; I'm not sure how much of the other code makes sense to re-use, as it's minimal. Let me know if you think otherwise.

leseb · 2025-05-19T12:55:54Z

llama_stack/templates/template.py

@@ -117,6 +118,9 @@ def run_config(
                __distro_dir__=f"~/.llama/distributions/{name}",
                db_name="registry.db",
            ),
+            inference_store=SqliteInferenceStoreConfig.sample_run_config(


Is it something we always want to enable?

I think this is good table-stake functionality that a users of Stack would reasonably expect and comes for free with OpenAI platform. Wdyt?

raghotham

Why have a separate inference_store? Cant we just reuse the kvstore that's already in the distro?

ehhuang · 2025-05-19T17:57:06Z

Why have a separate inference_store? Cant we just reuse the kvstore that's already in the distro?

My understanding is that kvstore doesn't support SQL-like queries? Are you suggesting to augment those interfaces?

ashwinb · 2025-05-19T19:08:12Z

Why have a separate inference_store? Cant we just reuse the kvstore that's already in the distro?

yeah agreed. I don't think we should introduce a specific inference store. We should add a "relational db store" or a "sql_store" (or some appopriate name) which allows us to target Relational DBs like Postgres in general.

that said, we should also see whether we truly need SQL queries here and how much can be done via decently designed schema around a KV store with ranges. (It is possible we truly need SQL support here.)

ehhuang · 2025-05-19T22:51:42Z

Why have a separate inference_store? Cant we just reuse the kvstore that's already in the distro?

yeah agreed. I don't think we should introduce a specific inference store. We should add a "relational db store" or a "sql_store" (or some appopriate name) which allows us to target Relational DBs like Postgres in general.

that said, we should also see whether we truly need SQL queries here and how much can be done via decently designed schema around a KV store with ranges. (It is possible we truly need SQL support here.)

Introduced a simple SqlStore API: llama_stack/providers/utils/sqlstore/api.py. PTAL

ashwinb · 2025-05-20T00:01:06Z

llama_stack/providers/utils/sqlstore/api.py

+    A protocol for a SQL store.
+    """
+
+    async def select(self, sql: str, params: tuple | None = None) -> list[dict]:


I think it would be useful to add a tables() for introspection also.

Is RelationalStore a better name? (I don't think so but putting it out there.)

I think SQL as an interface is way too opaque and later, way too difficult to optimize should the need be. We shouldn't need to parse the SQL query (and / or find that one needs to sanitize it and such). Instead I think we should do something like this:

def insert(self, table: str, row: Mapping[str, Any]) -> Any: ... def fetch_one(self, table: str, primary_key: Any) -> Mapping[str, Any] | None: ... def fetch_all( self, table: str, where: Mapping[str, Any] | str | None = None, *, limit: int | None = None, offset: int = 0, ) -> list[Mapping[str, Any]]: ... def update( self, table: str, where: Mapping[str, Any] | str, values: Mapping[str, Any] ) -> int: ... def delete(self, table: str, where: Mapping[str, Any] | str) -> int: ...

The most important aspect is the first argument -- a table. Every query should be within the context of a table. We can keep or add the execute power only if we start needing joins.

Alternately, what happened with adopting the sqlalchemy API for this?

I definitely think opaque sql should be a no-no.

I was trying to keep the expressivity, but this makes more sense. Updated the commit.

There's now a sqlalchemy sql store that's configurable by engine type.

Is RelationalStore a better name? (I don't think so but putting it out there.)

Don't have a strong opinion here either, can change if preferred.

# What does this PR do? ## Test Plan

ashwinb · 2025-05-20T22:39:22Z

pyproject.toml

@@ -41,6 +41,7 @@ dependencies = [
    "pillow",
    "h11>=0.16.0",
    "kubernetes",
+    "sqlalchemy[asyncio]>=2.0.41",


should we only add this dependency if a distro actually includes a sqlalchemy store?

ehhuang · 2025-05-20T22:39:27Z

requirements.txt

@@ -61,6 +62,7 @@ rsa==4.9
 setuptools==75.8.0
 six==1.17.0
 sniffio==1.3.1
+sqlalchemy==2.0.41


I think I need to move this to be part of build script?

facebook-github-bot added the CLA Signed This label is managed by the Meta Open Source bot. label May 19, 2025

ehhuang force-pushed the pr2200 branch 2 times, most recently from ac9f40f to 096cd0c Compare May 19, 2025 05:24

ehhuang changed the title ~~impl~~ feat: implement get chat completions APIs May 19, 2025

ehhuang force-pushed the pr2200 branch from 096cd0c to c2eb602 Compare May 19, 2025 05:35

ehhuang marked this pull request as ready for review May 19, 2025 05:38

ehhuang requested review from ashwinb, yanxi0830, hardikjshah, raghotham, terrytangyuan and leseb as code owners May 19, 2025 05:38

leseb reviewed May 19, 2025

View reviewed changes

raghotham reviewed May 19, 2025

View reviewed changes

ehhuang force-pushed the pr2200 branch from c2eb602 to 461d5f4 Compare May 19, 2025 18:10

ehhuang force-pushed the pr2200 branch from 461d5f4 to f9e6373 Compare May 19, 2025 22:48

ashwinb reviewed May 20, 2025

View reviewed changes

ehhuang force-pushed the pr2200 branch from f9e6373 to d6267dd Compare May 20, 2025 19:28

ehhuang requested a review from bbrowning as a code owner May 20, 2025 19:28

ehhuang force-pushed the pr2200 branch 3 times, most recently from c2d90bc to 9fb8299 Compare May 20, 2025 21:47

impl

f2e1fa3

# What does this PR do? ## Test Plan

ehhuang force-pushed the pr2200 branch from 9fb8299 to f2e1fa3 Compare May 20, 2025 22:01

ashwinb reviewed May 20, 2025

View reviewed changes

ehhuang commented May 20, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: implement get chat completions APIs #2200

feat: implement get chat completions APIs #2200

ehhuang commented May 19, 2025 •

edited

Loading

leseb May 19, 2025

leseb May 19, 2025

ehhuang May 19, 2025

leseb May 19, 2025

leseb May 19, 2025

leseb May 19, 2025

ehhuang May 19, 2025

leseb May 19, 2025

ehhuang May 19, 2025

raghotham left a comment

ehhuang commented May 19, 2025

ashwinb commented May 19, 2025

ehhuang commented May 19, 2025

ashwinb May 20, 2025

ashwinb May 20, 2025

ashwinb May 20, 2025

ehhuang May 20, 2025 •

edited

Loading

ashwinb May 20, 2025

ehhuang May 20, 2025

	if api.value == "inference" and run_config.inference_store:
	if api.value == Api.inference and run_config.inference_store:

feat: implement get chat completions APIs #2200

Are you sure you want to change the base?

feat: implement get chat completions APIs #2200

Conversation

ehhuang commented May 19, 2025 • edited Loading

What does this PR do?

Test Plan

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

raghotham left a comment

Choose a reason for hiding this comment

ehhuang commented May 19, 2025

ashwinb commented May 19, 2025

ehhuang commented May 19, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ehhuang May 20, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ehhuang commented May 19, 2025 •

edited

Loading

ehhuang May 20, 2025 •

edited

Loading