Various doc updates and expose some internals (#236)

DanielRosenwasser · web-flow · commit a1930e4231df · 2024-04-17T16:43:41.000-07:00
* Switch `sentiment` example to be a dataclass.

* Make it possible to actually reference the attributes of the internally-constructed models.

* Update doc comment.
diff --git a/python/examples/sentiment/demo.py b/python/examples/sentiment/demo.py
@@ -17,7 +17,7 @@ async def request_handler(message: str):
             print(result.message)
         else:
             result = result.value
-            print(f"The sentiment is {result['sentiment']}")
+            print(f"The sentiment is {result.sentiment}")
 
     file_path = sys.argv[1] if len(sys.argv) == 2 else None
     await process_requests("😀> ", file_path, request_handler)
diff --git a/python/examples/sentiment/schema.py b/python/examples/sentiment/schema.py
@@ -1,7 +1,8 @@
-from typing_extensions import Literal, TypedDict, Annotated, Doc
+from dataclasses import dataclass
+from typing_extensions import Literal, Annotated, Doc
 
-
-class Sentiment(TypedDict):
+@dataclass
+class Sentiment:
     """
     The following is a schema definition for determining the sentiment of a some user input.
     """
diff --git a/python/src/typechat/_internal/model.py b/python/src/typechat/_internal/model.py
@@ -39,10 +39,14 @@ class HttpxLanguageModel(TypeChatLanguageModel, AsyncContextManager):
     url: str
     headers: dict[str, str]
     default_params: dict[str, str]
+    # Specifies the maximum number of retry attempts.
+    max_retry_attempts: int = 3
+    # Specifies the delay before retrying in milliseconds.
+    retry_pause_seconds: float = 1.0
+    # Specifies how long a request should wait in seconds
+    # before timing out with a Failure.
+    timeout_seconds = 10
     _async_client: httpx.AsyncClient
-    _max_retry_attempts: int = 3
-    _retry_pause_seconds: float = 1.0
-    _timeout_seconds = 10
 
     def __init__(self, url: str, headers: dict[str, str], default_params: dict[str, str]):
         super().__init__()
@@ -74,7 +78,7 @@ async def complete(self, prompt: str | list[PromptSection]) -> Success[str] | Fa
                     self.url,
                     headers=headers,
                     json=body,
-                    timeout=self._timeout_seconds
+                    timeout=self.timeout_seconds
                 )
                 if response.is_success:
                     json_result = cast(
@@ -83,13 +87,13 @@ async def complete(self, prompt: str | list[PromptSection]) -> Success[str] | Fa
                     )
                     return Success(json_result["choices"][0]["message"]["content"] or "")
 
-                if response.status_code not in _TRANSIENT_ERROR_CODES or retry_count >= self._max_retry_attempts:
+                if response.status_code not in _TRANSIENT_ERROR_CODES or retry_count >= self.max_retry_attempts:
                     return Failure(f"REST API error {response.status_code}: {response.reason_phrase}")
             except Exception as e:
-                if retry_count >= self._max_retry_attempts:
+                if retry_count >= self.max_retry_attempts:
                     return Failure(str(e) or f"{repr(e)} raised from within internal TypeChat language model.")
 
-            await asyncio.sleep(self._retry_pause_seconds)
+            await asyncio.sleep(self.retry_pause_seconds)
             retry_count += 1
 
     @override
@@ -106,7 +110,7 @@ def __del__(self):
         except Exception:
             pass
 
-def create_language_model(vals: dict[str, str | None]) -> TypeChatLanguageModel:
+def create_language_model(vals: dict[str, str | None]) -> HttpxLanguageModel:
     """
     Creates a language model encapsulation of an OpenAI or Azure OpenAI REST API endpoint
     chosen by a dictionary of variables (typically just `os.environ`).
@@ -144,7 +148,7 @@ def required_var(name: str) -> str:
     else:
         raise ValueError("Missing environment variables for OPENAI_API_KEY or AZURE_OPENAI_API_KEY.")
 
-def create_openai_language_model(api_key: str, model: str, endpoint: str = "https://api.openai.com/v1/chat/completions", org: str = ""):
+def create_openai_language_model(api_key: str, model: str, endpoint: str = "https://api.openai.com/v1/chat/completions", org: str = "") -> HttpxLanguageModel:
     """
     Creates a language model encapsulation of an OpenAI REST API endpoint.
 
@@ -163,7 +167,7 @@ def create_openai_language_model(api_key: str, model: str, endpoint: str = "http
     }
     return HttpxLanguageModel(url=endpoint, headers=headers, default_params=default_params)
 
-def create_azure_openai_language_model(api_key: str, endpoint: str):
+def create_azure_openai_language_model(api_key: str, endpoint: str) -> HttpxLanguageModel:
     """
     Creates a language model encapsulation of an Azure OpenAI REST API endpoint.
 
diff --git a/python/src/typechat/_internal/validator.py b/python/src/typechat/_internal/validator.py
@@ -10,7 +10,7 @@
 
 class TypeChatValidator(Generic[T]):
     """
-    Validates JSON text against a given Python type.
+    Validates an object against a given Python type.
     """
 
     _adapted_type: pydantic.TypeAdapter[T]
