phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-06-28 19:04:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: support postgresql inference store (#2310)

Browse source 
# What does this PR do?
* Added support postgresql inference store
* Added 'oracle' template that demos how to config postgresql stores
(except for telemetry, which is not supported currently)


## Test Plan

llama stack build --template oracle --image-type conda --run
LLAMA_STACK_CONFIG=http://localhost:8321 pytest -s -v tests/integration/
--text-model accounts/fireworks/models/llama-v3p3-70b-instruct -k
'inference_store'
This commit is contained in:
ehhuang 2025-05-29 14:33:09 -07:00 • committed by Sumit Jaiswal
parent 6af13bbbf0
commit fbc8fc6eb5
No known key found for this signature in database
GPG key ID: A4604B39D64D6AEC
32 changed files with 516 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

38
llama_stack/providers/utils/sqlstore/sqlstore.py
View file

@ -5,6 +5,7 @@
# the root directory of this source tree.
from abc import abstractmethod
from enum import Enum
from pathlib import Path
from typing import Annotated, Literal
@ -21,7 +22,18 @@ class SqlStoreType(Enum):
postgres = "postgres"
class SqliteSqlStoreConfig(BaseModel):
class SqlAlchemySqlStoreConfig(BaseModel):
@property
@abstractmethod
def engine_str(self) -> str: ...
# TODO: move this when we have a better way to specify dependencies with internal APIs
@property
def pip_packages(self) -> list[str]:
return ["sqlalchemy[asyncio]"]
class SqliteSqlStoreConfig(SqlAlchemySqlStoreConfig):
type: Literal["sqlite"] = SqlStoreType.sqlite.value
db_path: str = Field(
default=(RUNTIME_BASE_DIR / "sqlstore.db").as_posix(),
@ -39,18 +51,26 @@ class SqliteSqlStoreConfig(BaseModel):
db_path="${env.SQLITE_STORE_DIR:" + __distro_dir__ + "}/" + db_name,
)
# TODO: move this when we have a better way to specify dependencies with internal APIs
@property
def pip_packages(self) -> list[str]:
return ["sqlalchemy[asyncio]"]
return super().pip_packages + ["aiosqlite"]
class PostgresSqlStoreConfig(BaseModel):
class PostgresSqlStoreConfig(SqlAlchemySqlStoreConfig):
type: Literal["postgres"] = SqlStoreType.postgres.value
host: str = "localhost"
port: str = "5432"
db: str = "llamastack"
user: str
password: str | None = None
@property
def engine_str(self) -> str:
return f"postgresql+asyncpg://{self.user}:{self.password}@{self.host}:{self.port}/{self.db}"
@property
def pip_packages(self) -> list[str]:
raise NotImplementedError("Postgres is not implemented yet")
return super().pip_packages + ["asyncpg"]
SqlStoreConfig = Annotated[
@ -60,12 +80,10 @@ SqlStoreConfig = Annotated[
def sqlstore_impl(config: SqlStoreConfig) -> SqlStore:
if config.type == SqlStoreType.sqlite.value:
from .sqlite.sqlite import SqliteSqlStoreImpl
if config.type in [SqlStoreType.sqlite.value, SqlStoreType.postgres.value]:
from .sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
impl = SqliteSqlStoreImpl(config)
elif config.type == SqlStoreType.postgres.value:
raise NotImplementedError("Postgres is not implemented yet")
impl = SqlAlchemySqlStoreImpl(config)
else:
raise ValueError(f"Unknown sqlstore type {config.type}")