DatabricksStore: PostgresStore Wrapper SDK #227
Conversation
jennsun
changed the title
jennsun
force-pushed
the
jenny-sun_data/j-lakebase-postgresstore2
branch
from
ea31ef1 to
272c62d
Compare
| Wrapper around LangGraph's PostgresStore that uses a Lakebase | ||
| connection pool and borrows a connection per call. | ||
| """ | ||
|
|
There was a problem hiding this comment.
note: instead of directly subclassing, I wrapped PostgresStore instead to simplify the agent authoring CUJ using the store/handle connections for the user
Previously:
with self.pool.connection() as conn:
store = PostgresStore(conn=conn)
results = store.search(namespace, query=query, limit=5)
if we subclassed it would look like:
with self.pool.connection() as conn:
store = DatabricksStore(conn=conn)
results = store.search(namespace, query=query, limit=5)
current
results = self.store.search(namespace, query=query, limit=5)
let me know if we prefer to subclass instead
There was a problem hiding this comment.
This approach looks good to me, hopefully advanced connection management isn't needed most of the time + it's good to hide the LakebasePool details
|
|
||
| class DatabricksStore: | ||
| """ | ||
| Wrapper around LangGraph's PostgresStore that uses a Lakebase |
There was a problem hiding this comment.
Can we document this instead for its purpose e.g. explain that this class provides APIs for working with long-term memory, extending LangGraph's Store interface?
| _store_imports_available = False | ||
|
|
||
|
|
||
| class DatabricksStore: |
There was a problem hiding this comment.
Is there a base Store class we should be subclassing?
There was a problem hiding this comment.
smurching
left a comment
smurching
left a comment
There was a problem hiding this comment.
Minor comments but mostly looks good!
| Currently delegates to sync batch() - for true async support, | ||
| would need async-compatible connection pooling. | ||
| """ | ||
| return self.batch(ops) |
There was a problem hiding this comment.
if we are going to provide standard functions like _normalize_ns_label and _user_namespace, would it make sense to put them here as well? just to reduce the total LoC in the example
There was a problem hiding this comment.
also a comment on the example: can we copy the langgraph example in using output_to_responses_items_stream https://docs.databricks.com/aws/en/generative-ai/agent-framework/author-agent?language=LangGraph#responsesagent-examples
There was a problem hiding this comment.
Yes to #1 for refactoring those functions related to the key-value storage -- for #2, I believe we're already doing this in the long-term example here - let me know if this is what you were thinking: https://github.com/databricks-eng/universe/pull/1443034/files?w=1#diff-fbb7ade56e8f48e42bc41e5a4672a52e301579fd4f7dc2f7cd0f1e5a30a0babbR463
There was a problem hiding this comment.
sounds good for #1 -- can we add some unit tests for the behavior we expect?
bbqiu
left a comment
bbqiu
left a comment
There was a problem hiding this comment.
overall looks great! just a few questions and want to see if there are more of the helper functions in the nb into db-ai-bridge as helper functions
bbqiu
left a comment
bbqiu
left a comment
There was a problem hiding this comment.
LGTM thank you for moving the lazy init logic into the agent rather than the SDK since we might want to reuse the SDK in envs other than model serving
Previous CUJ with PostgresStore:
with new DatabricksStore:
Used for long-term user preference store,
long-term memory notebook: https://eng-ml-liteswap-us-east-1-2.staging.cloud.databricks.com/editor/notebooks/2695522312070529?o=1036778114150303#command/8820576549438873
example endpoint: https://eng-ml-liteswap-us-east-1-2.staging.cloud.databricks.com/ml/endpoints/agents_main-default-j-longterm-sdk-ns?o=1036778114150303