Repository: MadcowD/ell

Issue:
Ollama example not working
The ollama example provided (examples/ollama_example.py) with the following error:

`AttributeError: module 'ell.models.ollama' has no attribute 'register_models'`



Discussion:

Patch:
diff --git a/examples/calculator_structured.py b/examples/calculator_structured.py
deleted file mode 100644
index bdd02b0..0000000
--- a/examples/calculator_structured.py
+++ /dev/null
@@ -1,73 +0,0 @@
-import dataclasses
-import ell
-from typing import Any, Literal, Type, Union
-import pydantic
-
-from ell.stores.sql import SQLiteStore
-
-
-
-
-@dataclasses.dataclass
-class Add:
-    op: Literal["+"]
-    a: float
-    b: float
-
-
-@dataclasses.dataclass
-class Sub:
-    op: Literal["-"]
-    a: float
-    b: float
-
-
-@dataclasses.dataclass
-class Mul:
-    op: Literal["*"]
-    a: float
-    b: float
-
-
-@dataclasses.dataclass
-class Div:
-    op: Literal["/"]
-    a: float
-    b: float
-
-
-CalcOp = Union[Add, Sub, Mul, Div]
-
-
-@ell.simple(model="gpt-4o", temperature=0.1)
-def parse_json(task: str, type: Type[Any]):
-    return [
-        ell.system(
-            f"""You are a JSON parser. You respond only in JSON. Do not format using markdown."""
-        ),
-        ell.user(
-            f"""You are given the following task: "{task}"
-            Parse the task into the following type:
-            {pydantic.TypeAdapter(type).json_schema()}
-            """
-        )
-    ]
-
-
-def calc_structured(task: str) -> float:
-    output = parse_json(task, CalcOp)
-    structured = pydantic.TypeAdapter(CalcOp).validate_json(output)
-    match structured.op:
-        case "+":
-            return structured.a + structured.b
-        case "-":
-            return structured.a - structured.b
-        case "*":
-            return structured.a * structured.b
-        case "/":
-            return structured.a / structured.b
-
-
-if __name__ == "__main__":
-    ell.init(store='./logdir', autocommit=True, verbose=True)
-    print(calc_structured("What is two plus two?"))
diff --git a/examples/ollama_example.py b/examples/ollama_example.py
deleted file mode 100644
index 4714c3c..0000000
--- a/examples/ollama_example.py
+++ /dev/null
@@ -1,10 +0,0 @@
-import ell
-
-@ell.simple(model="llama3", temperature=0.1)
-def write_a_story():
-    return "write me a story"
-
-
-ell.models.ollama.register_models(api_base="http://localhost:11434")
-
-write_a_story()
\ No newline at end of file
diff --git a/examples/providers/anthropic_ex.py b/examples/providers/anthropic_ex.py
new file mode 100644
index 0000000..1e1c80b
--- /dev/null
+++ b/examples/providers/anthropic_ex.py
@@ -0,0 +1,26 @@
+""" 
+Anthropic example: pip install ell-ai[anthropic]
+"""
+import ell
+import anthropic
+
+ell.init(verbose=True)
+
+# custom client
+client = anthropic.Anthropic()
+
+@ell.simple(model='claude-3-5-sonnet-20240620', client=client, max_tokens=10)
+def chat(prompt: str) -> str:
+    return prompt
+
+print(chat("Hello, how are you?"))
+
+# Models are automatically registered!
+@ell.simple(model='claude-3-5-sonnet-20240620', max_tokens=10)
+def use_default_client(prompt: str) -> str:
+    return prompt
+
+print(use_default_client("Hello, how are you?"))
+
+
+
diff --git a/examples/providers/groq_ex.py b/examples/providers/groq_ex.py
new file mode 100644
index 0000000..a0a527c
--- /dev/null
+++ b/examples/providers/groq_ex.py
@@ -0,0 +1,26 @@
+"""
+Groq example: pip install ell-ai[groq]
+"""
+import ell
+import groq
+
+
+ell.init(verbose=True, store='./logdir')
+ell.models.groq.register() # use GROQ_API_KEY env var
+# ell.models.groq.register(api_key="gsk-") # 
+
+@ell.simple(model="llama3-8b-8192", temperature=0.1)
+def write_a_story(about : str):
+    return f"write me a story about {about}"
+
+write_a_story("cats")
+
+# or use the client directly
+client = groq.Groq()
+
+@ell.simple(model="llama3-8b-8192", temperature=0.1, client=client)
+def write_a_story_with_client(about : str):
+    return f"write me a story about {about}"
+
+write_a_story_with_client("cats")
+
diff --git a/examples/providers/ollama_ex.py b/examples/providers/ollama_ex.py
new file mode 100644
index 0000000..f8bd5fa
--- /dev/null
+++ b/examples/providers/ollama_ex.py
@@ -0,0 +1,20 @@
+"""
+Ollama example.
+"""
+import ell
+import openai
+
+# Use models automatically registered by asking ollama 
+ell.models.ollama.register(api_base="http://localhost:11434/v1")
+@ell.simple(model="llama3", temperature=0.1)
+def write_a_story():
+    return "write me a story"
+
+
+# Or use the client firectly
+client = openai.Client(
+    base_url="http://localhost:11434/v1", api_key="ollama"  # required but not used
+)
+@ell.simple(model="llama3", temperature=0.1, max_tokens=10, client=client)
+def write_a_story_with_client():
+    return "write me a short"
diff --git a/examples/providers/openai_ex.py b/examples/providers/openai_ex.py
new file mode 100644
index 0000000..9225c63
--- /dev/null
+++ b/examples/providers/openai_ex.py
@@ -0,0 +1,23 @@
+import ell
+import openai
+
+ell.init(verbose=True)
+
+# custom client
+client = openai.Client()
+
+@ell.simple(model='gpt-4o', client=client)
+def chat(prompt: str) -> str:
+    return prompt
+
+print(chat("Hello, how are you?"))
+
+# Models are automatically registered!
+@ell.simple(model='gpt-4o')
+def use_default_client(prompt: str) -> str:
+    return prompt
+
+print(use_default_client("Hello, how are you?"))
+
+
+
diff --git a/examples/providers/vllm_ex.py b/examples/providers/vllm_ex.py
new file mode 100644
index 0000000..e733008
--- /dev/null
+++ b/examples/providers/vllm_ex.py
@@ -0,0 +1,25 @@
+"""
+vLLM example.
+"""
+from openai import OpenAI
+import ell
+# vllm serve NousResearch/Meta-Llama-3-8B-Instruct --dtype auto --api-key token-abc123
+
+client = OpenAI(
+    base_url="http://localhost:8000/v1",
+    api_key="token-abc123",
+)
+
+@ell.simple(model="NousResearch/Meta-Llama-3-8B-Instruct", client=client, temperature=0.1)
+def write_a_story(about : str):
+    return f"write me a story about {about}"
+
+# or register models
+ell.config.register_model("NousResearch/Meta-Llama-3-8B-Instruct", client)
+
+# no need to specify client!
+@ell.simple(model="NousResearch/Meta-Llama-3-8B-Instruct", temperature=0.1)
+def write_a_story_no_client(about : str):
+    return f"write me a story about {about}"
+
+write_a_story()
\ No newline at end of file
diff --git a/poetry.lock b/poetry.lock
index a4dd298..4aa3072 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -578,6 +578,25 @@ files = [
 docs = ["Sphinx", "furo"]
 test = ["objgraph", "psutil"]
 
+[[package]]
+name = "groq"
+version = "0.11.0"
+description = "The official Python library for the groq API"
+optional = true
+python-versions = ">=3.7"
+files = [
+    {file = "groq-0.11.0-py3-none-any.whl", hash = "sha256:e328531c979542e563668c62260aec13b43a6ee0ca9e2fb22dff1d26f8c8ce54"},
+    {file = "groq-0.11.0.tar.gz", hash = "sha256:dbb9aefedf388ddd4801ec7bf3eba7f5edb67948fec0cd2829d97244059f42a7"},
+]
+
+[package.dependencies]
+anyio = ">=3.5.0,<5"
+distro = ">=1.7.0,<2"
+httpx = ">=0.23.0,<1"
+pydantic = ">=1.9.0,<3"
+sniffio = "*"
+typing-extensions = ">=4.7,<5"
+
 [[package]]
 name = "h11"
 version = "0.14.0"
@@ -2292,11 +2311,12 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
 type = ["pytest-mypy"]
 
 [extras]
-all = ["anthropic", "openai"]
+all = ["anthropic", "groq", "openai"]
 anthropic = ["anthropic"]
+groq = ["groq"]
 openai = ["openai"]
 
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9"
-content-hash = "1db008749153abc0dba7bcacc6b92bb3d15d47e10addbec9c175a8981a489b3a"
+content-hash = "621457432d7ef7d5348d1909b7dcc4feec959a7a46d9101687e4f76fd3e9dff8"
diff --git a/pyproject.toml b/pyproject.toml
index 336e907..4d691ec 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -33,6 +33,7 @@ colorama = "^0.4.6"
 cattrs = "^23.2.3"
 openai = "^1.47.0"
 anthropic = { version = "^0.34.2", optional = true }
+groq = { version = "^0.11.0", optional = true }
 sqlmodel = "^0.0.21"
 uvicorn = "^0.30.3"
 requests = "^2.32.3"
@@ -49,7 +50,8 @@ sphinx-rtd-theme = "^2.0.0"
 [tool.poetry.extras]
 openai = ["openai"]
 anthropic = ["anthropic"]
-all = ["openai", "anthropic"]
+groq = ["groq"]
+all = ["openai", "anthropic", "groq"]
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
diff --git a/src/ell/configurator.py b/src/ell/configurator.py
index 7f1ef62..a0758f5 100644
--- a/src/ell/configurator.py
+++ b/src/ell/configurator.py
@@ -187,9 +187,8 @@ def get_store() -> Union[Store, None]:
 # Will be deprecated at 0.1.0 
 
 # You can add more helper functions here if needed
-@wraps(config.register_provider)
-def register_provider(*args, **kwargs) -> None:
-    return config.register_provider(*args, **kwargs)
+def register_provider(provider: Provider, client_type: Type[Any]) -> None:
+    return config.register_provider(provider, client_type)
 
 # Deprecated now (remove at 0.1.0)
 def set_store(*args, **kwargs) -> None:
diff --git a/src/ell/models/__init__.py b/src/ell/models/__init__.py
index b46d8a0..865000b 100644
--- a/src/ell/models/__init__.py
+++ b/src/ell/models/__init__.py
@@ -8,4 +8,5 @@ For example, to register an OpenAI model:
 
 import ell.models.openai
 import ell.models.anthropic
-import ell.models.ollama
\ No newline at end of file
+import ell.models.ollama
+import ell.models.groq
\ No newline at end of file
diff --git a/src/ell/models/groq.py b/src/ell/models/groq.py
new file mode 100644
index 0000000..25d7cbb
--- /dev/null
+++ b/src/ell/models/groq.py
@@ -0,0 +1,12 @@
+from typing import Optional
+from ell.configurator import config
+
+try:
+    from groq import Groq
+    def register(client: Optional[Groq] = None, **client_kwargs):
+        if client is None:
+            client = Groq(**client_kwargs)
+        for model in client.models.list().data:
+            config.register_model(model.id, default_client=client, supports_streaming=True)
+except ImportError:
+    pass
\ No newline at end of file
diff --git a/src/ell/models/ollama.py b/src/ell/models/ollama.py
index 7115c41..87bf47c 100644
--- a/src/ell/models/ollama.py
+++ b/src/ell/models/ollama.py
@@ -3,6 +3,7 @@ import openai
 import requests
 import logging
 
+#XXX: May be deprecated soon because of the new provider framework.
 logger = logging.getLogger(__name__)
 client = None
 
diff --git a/src/ell/providers/__init__.py b/src/ell/providers/__init__.py
index 763dfc0..0f24e7e 100644
--- a/src/ell/providers/__init__.py
+++ b/src/ell/providers/__init__.py
@@ -1,6 +1,6 @@
 import ell.providers.openai
+import ell.providers.groq
 import ell.providers.anthropic
-# import ell.providers.groq
 # import ell.providers.mistral
 # import ell.providers.cohere
 # import ell.providers.gemini
diff --git a/src/ell/providers/anthropic.py b/src/ell/providers/anthropic.py
index 4d99fa8..5f0e025 100644
--- a/src/ell/providers/anthropic.py
+++ b/src/ell/providers/anthropic.py
@@ -148,10 +148,11 @@ try:
             metadata["usage"] = usage
             return tracked_results, metadata
 
-    # XXX: Could register a true base class.
-    register_provider(AnthropicProvider(), anthropic.Anthropic)
-    register_provider(AnthropicProvider(), anthropic.AnthropicBedrock)
-    register_provider(AnthropicProvider(), anthropic.AnthropicVertex)
+    # XXX: Make a singleton.
+    anthropic_provider = AnthropicProvider()
+    register_provider(anthropic_provider, anthropic.Anthropic)
+    register_provider(anthropic_provider, anthropic.AnthropicBedrock)
+    register_provider(anthropic_provider, anthropic.AnthropicVertex)
 
 except ImportError:
     raise
diff --git a/src/ell/providers/groq.py b/src/ell/providers/groq.py
new file mode 100644
index 0000000..de0ae53
--- /dev/null
+++ b/src/ell/providers/groq.py
@@ -0,0 +1,26 @@
+"""
+Groq provider.
+"""
+
+from ell.providers.openai import OpenAIProvider
+from ell.configurator import register_provider
+
+
+try:
+    import groq
+    class GroqProvider(OpenAIProvider):
+        dangerous_disable_validation = True
+        def translate_to_provider(self, *args, **kwargs):
+            params = super().translate_to_provider(*args, **kwargs)
+            params.pop('stream_options', None)
+            return params
+        
+        def translate_from_provider(self, *args, **kwargs):
+            res, meta = super().translate_from_provider(*args, **kwargs)
+            if not meta['usage']:
+                meta['usage'] = meta['x_groq']['usage']
+            return res, meta
+    register_provider(GroqProvider(), groq.Client)
+except ImportError:
+    pass
+
diff --git a/src/ell/providers/openai.py b/src/ell/providers/openai.py
index c79b322..e8b17be 100644
--- a/src/ell/providers/openai.py
+++ b/src/ell/providers/openai.py
@@ -107,15 +107,14 @@ try:
                 message_streams = defaultdict(list)
                 role : Optional[str] = None
                 for chunk in stream: 
-        
-                    if hasattr(chunk, "usage") and chunk.usage: metadata.update(chunk.model_dump(exclude={"choices"})) 
+                    metadata.update(chunk.model_dump(exclude={"choices"})) 
                     
                     for chat_compl_chunk in chunk.choices:
                         message_streams[chat_compl_chunk.index].append(chat_compl_chunk)
                         delta = chat_compl_chunk.delta
                         role = role or delta.role
                         if  chat_compl_chunk.index == 0 and logger:
-                            logger(delta.content, is_refusal=delta.refusal)
+                            logger(delta.content, is_refusal=hasattr(delta, "refusal") and delta.refusal)
                 for _, message_stream in sorted(message_streams.items(), key=lambda x: x[0]):
                     text = "".join((choice.delta.content or "") for choice in message_stream)
                     messages.append(
@@ -128,7 +127,7 @@ try:
                 for oai_choice in chat_completion.choices: 
                     role = oai_choice.message.role
                     content_blocks = []
-                    if (refusal := (message := oai_choice.message).refusal):
+                    if (hasattr(message := oai_choice.message, "refusal") and (refusal := message.refusal)):
                         raise ValueError(refusal)
                     if hasattr(message, "parsed"):
                         if (parsed := message.parsed): 
@@ -156,6 +155,7 @@ try:
             return messages, metadata
 
 
+    # xx: singleton needed
     openai_provider = OpenAIProvider()
     register_provider(openai_provider, openai.Client)
 except ImportError:


Test patch:
diff --git a/tests/test_openai_provider.py b/tests/test_openai_provider.py
index 2ecd01b..d7ec48c 100644
--- a/tests/test_openai_provider.py
+++ b/tests/test_openai_provider.py
@@ -190,7 +190,10 @@ class TestOpenAIProvider:
             assert messages == [
                 Message(role="assistant", content=[ContentBlock(text="Hello")])
             ]
-            assert metadata == {}
+            assert metadata['id'] == 'chatcmpl-123'
+            assert not metadata['usage']
+            assert metadata['model'] == 'gpt-4'
+            assert metadata['object'] == 'chat.completion.chunk'
 
     def test_translate_from_provider_non_streaming(self, provider, ell_call_params):
         provider_call_params = {"stream": False}
@@ -355,7 +358,7 @@ class TestOpenAIProvider:
         assert messages == [
             Message(role="assistant", content=[ContentBlock(text="Hello World")])
         ]
-        assert metadata == {}
+
 
     # Suggested Test for _content_block_to_openai_format
 def test_content_block_to_openai_format():