Skip to content

Commit

Permalink
fix: #4 supports older pydantic as well
Browse files Browse the repository at this point in the history
  • Loading branch information
AlmogBaku committed Mar 14, 2024
1 parent 07b6ee2 commit 0ffda1c
Show file tree
Hide file tree
Showing 4 changed files with 10 additions and 7 deletions.
9 changes: 6 additions & 3 deletions openai_streaming/decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from typing import get_args

from docstring_parser import parse
from openai.types.beta.assistant import ToolFunction
from openai.types.beta import FunctionTool
from openai.types.shared import FunctionDefinition
from pydantic import create_model

Expand Down Expand Up @@ -72,14 +72,17 @@ async def error_message(typ: str, description: AsyncGenerator[str, None]):
docstring = parse(func.__doc__ or "")

# prepare the parameters(arguments)
parameters = model.model_json_schema()
try:
parameters = model.model_json_schema()
except Exception as e:
parameters = model.schema() # Fallback to the default schema

# extract parameter documentations from the docstring
for param in docstring.params:
if (name := param.arg_name) in parameters["properties"] and (description := param.description):
parameters["properties"][name]["description"] = description

func.openai_schema = ToolFunction(type='function', function=FunctionDefinition(
func.openai_schema = FunctionTool(type='function', function=FunctionDefinition(
name=func.__name__,
description=docstring.short_description,
parameters=parameters,
Expand Down
2 changes: 1 addition & 1 deletion openai_streaming/stream_processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,7 +222,7 @@ def _process_message(
:return: Generator
"""
choice = message.choices[0]
if not choice.model_fields.get("delta"):
if not hasattr(choice, "delta"):
raise LookupError("No delta in choice")

delta = message.choices[0].delta
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ license = {text = "MIT"}
readme = "README.md"
keywords = ["openai", "gpt", "llm", "streaming", "stream", "generator"]
dependencies= [
"openai>=1.0.0,<2.0.0",
"openai>=1.14.0,<2.0.0",
"json-streamer>=0.1.0,<0.2.0",
"pydantic>=2.0.2,<3.0.0",
"docstring-parser>=0.15"
Expand Down
4 changes: 2 additions & 2 deletions requirements.txt
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
openai==1.11.1
openai==1.14.0
json-streamer==0.1.0
pydantic==2.6.1
pydantic==2.6.4
docstring-parser==0.15

0 comments on commit 0ffda1c

Please sign in to comment.