fix

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

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,

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
+"""