diff --git a/examples/sqlmodel-integration/patch_sql.py b/examples/sqlmodel-integration/patch_sql.py index 5fab57700..c5ed34733 100644 --- a/examples/sqlmodel-integration/patch_sql.py +++ b/examples/sqlmodel-integration/patch_sql.py @@ -1,3 +1,6 @@ +""" +TODO: add the created_at, id and called it completion vs response +""" try: import importlib @@ -14,9 +17,9 @@ from functools import wraps import openai import inspect -from typing import Callable, Optional +from typing import Callable -from sa import engine, ChatCompletion, Message +from sa import ChatCompletion, Message # Check if the engine is asynchronous @@ -40,6 +43,7 @@ async def async_insert_chat_completion( async with AsyncSessionLocal() as session: chat = ChatCompletion( + id=kwargs.pop("id", None), messages=[ Message(**message, index=ii) for (ii, message) in enumerate(messages) ], @@ -59,10 +63,11 @@ def sync_insert_chat_completion( ): with Session(engine) as session: chat = ChatCompletion( + id=kwargs.pop("id", None), messages=[ Message(**message, index=ii) for (ii, message) in enumerate(messages) ], - responses=[Message(**response) for response in responses], + responses=[Message(index=response["index"], **response["message"], is_response=True) for response in responses], **kwargs, ) session.add(chat) @@ -91,8 +96,9 @@ def add_sql_alchemy(func: Callable) -> Callable: async def new_chatcompletion(*args, **kwargs): # type: ignore response = await func(*args, **kwargs) save_chat_completion( - messages=response.messages, + messages=kwargs.pop("messages", []), responses=response.choices, + id=response["id"], **response["usage"], **kwargs, ) @@ -104,11 +110,13 @@ async def new_chatcompletion(*args, **kwargs): # type: ignore def new_chatcompletion(*args, **kwargs): response = func(*args, **kwargs) save_chat_completion( - messages=response.messages, + messages=kwargs.pop("messages", []), responses=response.choices, + id=response["id"], **response["usage"], **kwargs, ) + response._completion_id = response["id"] return response return new_chatcompletion diff --git a/examples/sqlmodel-integration/sa.py b/examples/sqlmodel-integration/sa.py index 3259359d3..65a4e2a53 100644 --- a/examples/sqlmodel-integration/sa.py +++ b/examples/sqlmodel-integration/sa.py @@ -16,6 +16,7 @@ class Message(Base): __tablename__ = "message" id = Column(Integer, primary_key=True, index=True) + index = Column(Integer) role = Column(String) content = Column(String, index=True) arguments = Column(String) @@ -58,10 +59,11 @@ class ChatCompletion(Base): if __name__ == "__main__": - sqlite_file_name = "example.db" + sqlite_file_name = "chat.db" sqlite_url = f"sqlite:///{sqlite_file_name}" engine = create_engine(sqlite_url, echo=True) + # create all the tables Base.metadata.create_all(engine) with Session(engine) as session: @@ -82,10 +84,9 @@ class ChatCompletion(Base): Message(role="user", content="My computer is not turning on."), ] chat_completion = ChatCompletion( - id="test", + id="test2", messages=messages, responses=responses, - finish_reason="stop", temperature=0.1, model="gpt-3.5-turbo-0613", max_tokens=1000, diff --git a/examples/sqlmodel-integration/test.py b/examples/sqlmodel-integration/test.py index 1cba32d3e..14645d63a 100644 --- a/examples/sqlmodel-integration/test.py +++ b/examples/sqlmodel-integration/test.py @@ -1,4 +1,9 @@ import openai +from sqlalchemy import create_engine +from patch_sql import instrument_with_sqlalchemy + +engine = create_engine("sqlite:///chat.db", echo=True) +instrument_with_sqlalchemy(engine) resp = openai.ChatCompletion.create( model="gpt-3.5-turbo-0613", @@ -9,4 +14,13 @@ } ], ) -print(resp) +""" +2023-08-19 14:57:04,339 INFO sqlalchemy.engine.Engine BEGIN (implicit) +2023-08-19 14:57:04,340 INFO sqlalchemy.engine.Engine INSERT INTO chatcompletion +2023-08-19 14:57:04,340 INFO sqlalchemy.engine.Engine [generated in 0.00021s] ('chatcmpl-7pOEKM94WoZX9IJ9rRzPacrylzlOR', '2023-08-19 21:57:04.340462', None, ... +2023-08-19 14:57:04,341 INFO sqlalchemy.engine.Engine INSERT INTO message +2023-08-19 14:57:04,342 INFO sqlalchemy.engine.Engine [generated in 0.00012s] +2023-08-19 14:57:04,342 INFO sqlalchemy.engine.Engine INSERT INTO message +2023-08-19 14:57:04,342 INFO sqlalchemy.engine.Engine [cached since 0.0004963s ago] +2023-08-19 14:57:04,342 INFO sqlalchemy.engine.Engine COMMIT +""" \ No newline at end of file