examples

jxnl · Aug 21, 2023 · 6f6e87f · 6f6e87f
1 parent 84ea095
commit 6f6e87f
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 20 deletions.
diff --git a/examples/sqlmodel-integration/patch_sql.py b/examples/sqlmodel-integration/patch_sql.py
@@ -1,19 +1,11 @@
-"""
-TODO: add the created_at, id and called it completion vs response
-"""
-
 try:
-    import importlib
-
-    importlib.import_module("sqlalchemy")
+    from sqlalchemy.orm import Session
 except ImportError:
     import warnings
 
-    warnings.warn("SQLAlchemy is not installed. Please install it to use this feature.")
-
-from sqlalchemy import Engine
-from sqlalchemy.orm import Session
+    warnings.warn("SQLAlchemy is not installed. Please install it to use this feature `pip install sqlalchemy`")
 
+import time
 from functools import wraps
 import openai
 import inspect
@@ -42,18 +34,19 @@ def sync_insert_chat_completion(
     responses: list[dict] = [],
     **kwargs,
 ):
-    with Session(engine) as session:
+    with Session(engine) as session: 
         chat = ChatCompletion(
             id=kwargs.pop("id", None),
             created_at=kwargs.pop("created", None),
             functions=json.dumps(kwargs.pop("functions", None)),
             function_call=json.dumps(kwargs.pop("function_call", None)),
+            latency_ms=kwargs.pop("latency_ms", None),
             messages=[
                 sql_message(index=ii, message=message)
                 for (ii, message) in enumerate(messages)
             ],
             responses=[
-                sql_message(index=resp["index"], message=resp.message, is_response=True)
+                sql_message(index=resp["index"], message=resp.message, is_response=True) # type: ignore
                 for resp in responses
             ],
             **kwargs,
@@ -69,11 +62,15 @@ def add_sql_alchemy(func: Callable) -> Callable:
 
             @wraps(func)
             async def new_chatcompletion(*args, **kwargs):  # type: ignore
+
+                start_ms = time.time()
                 response = await func(*args, **kwargs)
+                latency_ms = round((time.time() - start_ms) * 1000)
                 sync_insert_chat_completion(
                     engine,
                     messages=kwargs.pop("messages", []),
                     responses=response.choices,
+                    latency_ms=latency_ms,
                     id=response["id"],
                     **response["usage"],
                     **kwargs,
@@ -84,17 +81,19 @@ async def new_chatcompletion(*args, **kwargs):  # type: ignore
 
             @wraps(func)
             def new_chatcompletion(*args, **kwargs):
+                start_ms = time.time()
                 response = func(*args, **kwargs)
+                latency_ms = round((time.time() - start_ms) * 1000)
 
                 sync_insert_chat_completion(
                     engine,
                     messages=kwargs.pop("messages", []),
                     responses=response.choices,
                     id=response["id"],
+                    latency_ms=latency_ms,
                     **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
@@ -60,6 +60,7 @@ class ChatCompletion(Base):
     prompt_tokens = Column(Integer)
     completion_tokens = Column(Integer)
     total_tokens = Column(Integer)
+    latency_ms = Column(Integer)
     functions = Column(String)  # TODO: make this a foreign key
     function_call = Column(String)  # TODO: make this a foreign key
 

diff --git a/examples/sqlmodel-integration/test.py b/examples/sqlmodel-integration/test.py
@@ -1,6 +1,6 @@
 import openai
 
-from patch_sql import instrument_chat_completion_sa, instrument_with_sa_engine
+from patch_sql import instrument_chat_completion_sa
 from instructor.patch import patch_chatcompletion_with_response_model
 
 from typing import Optional
@@ -11,7 +11,7 @@
 # SQLAlchemy allows you to support any database!
 engine = create_engine("sqlite:///openai.db", echo=True)
 
-instrument_with_sa_engine(engine)
+instrument_chat_completion_sa(engine)
 
 # This adds a response_model parameter to the ChatCompletion.careate method
 patch_chatcompletion_with_response_model()
@@ -26,13 +26,11 @@ class UserDetails(SQLModel, table=True):
 
     def save(self):
         with Session(engine) as session:
-
-            if hasattr(self, "_raw_response"):
-                self.completion_id = self._raw_response["id"]
-
             session.add(self)
             session.commit()
 
+SQLModel.metadata.create_all(engine)
+
 
 resp: UserDetails = openai.ChatCompletion.create(
     model="gpt-3.5-turbo-0613",