Merge pull request #573 from PrefectHQ/docs-+-ai_fn-QOL

docs fix + QOL parsing stuff (rm parse raw for json.loads)
PrefectHQ · Sep 12, 2023 · 620c637 · 620c637
2 parents 14acab5 + f4e3de6
commit 620c637
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 17 deletions.
diff --git a/docs/examples/classification_api.md b/docs/examples/classification_api.md
@@ -23,20 +23,14 @@ In this example, we'll show how to
 
     ```python
     from marvin import ai_fn, settings
-    from enum import Enum
-
-    class Category(Enum):
-        SPORTS = "sports"
-        POLITICS = "politics"
-        TECHNOLOGY = "technology"
-        # Add more categories as needed
+    from typing import Literal
 
     settings.openai.api_key = 'API_KEY' 
 
     @ai_fn
-    def classify_text(text: str) -> Category:
+    def classify_text(text: str) -> Literal['sports', 'politics', 'technology']:
         '''
-            Classifies the passed `text` into one of the predefined categories. 
+            Correctly classifies the passed `text` into one of the predefined categories. 
         '''
 
     ```
@@ -61,16 +55,19 @@ In the following example, we will demonstrate how to deploy the AI function we j
     ```python
     from fastapi import FastAPI
     from marvin import ai_fn, settings
+    from typing import Literal
+
+    settings.openai.api_key = 'API_KEY' 
 
     app = FastAPI()
 
     settings.openai.api_key = 'API_KEY'
 
     @app.post("/classify_text/")
     @ai_fn
-    def classify_text(text: str) -> Category:
+    def classify_text(text: str) -> Literal['sports', 'politics', 'technology']:
         '''
-            Classifies the passed `text` into one of the predefined categories. 
+            Correctly classifies the passed `text` into one of the predefined categories. 
         '''
     ```
 

diff --git a/docs/examples/extraction_api.md b/docs/examples/extraction_api.md
@@ -24,13 +24,14 @@ In this example, we'll show how to
 
     ```python
     from marvin import ai_fn, settings
+    from typing import Any
 
     settings.openai.api_key = 'API_KEY' 
 
     @ai_fn
-    def extract_person(text: str) -> dict[str, Any]:
+    def extract_person(text: str) -> dict:
         '''
-            Extracts a persons `birth_year`, `first_name` and `last_name`
+            Correctly infers a persons `birth_year`, `first_name` and `last_name`
             from the passed `text`. 
         '''
 
@@ -63,9 +64,9 @@ In the following example, we will demonstrate how to deploy the AI function we j
 
     @app.post("/extract_person/")
     @ai_fn
-    def extract_person(text: str) -> dict[str, Any]:
+    def extract_person(text: str) -> dict:
         '''
-            Extracts a persons `birth_year`, `first_name` and `last_name`
+            Correctly infers a persons `birth_year`, `first_name` and `last_name`
             from the passed `text`. 
         '''
         

diff --git a/src/marvin/core/ChatCompletion/base/__init__.py b/src/marvin/core/ChatCompletion/base/__init__.py
@@ -203,8 +203,10 @@ def call_function(self, as_message=True):
         """
         name, raw_arguments = itemgetter("name", "arguments")(self.function_call())
         function = self.callable_registry.get(name)
-        arguments = function.model.parse_raw(raw_arguments)
-        value = function(**arguments.dict(exclude_none=True))
+        try:
+            value = function(**json.loads(raw_arguments))
+        except Exception:
+            value = function(data=json.loads(raw_arguments))
         if as_message:
             return {"role": "function", "name": name, "content": str(value)}
         else: