cohere example

Author: alizenhom

src/examples/litellm_example/basic.py

@@ -13,23 +13,29 @@
 
 
 @with_langtrace_root_span("Litellm Example OpenAI")
-def openAI():
+def openAI(streaming=False):
     response = completion(
         model="gpt-3.5-turbo",
         messages=[
             {"content": "respond only in Yoda speak.", "role": "system"},
             {"content": "Hello, how are you?", "role": "user"},
         ],
+        stream=streaming,
+        stream_options={"include_usage": True},
     )
-    return response
+    if streaming:
+        for _ in response:
+            pass
+    else:
+        return response
 
 
 # @with_langtrace_root_span("Litellm Example Anthropic Completion")
 def anthropic(streaming=False):
     try:
 
         response = completion(
-            model="claude-2",
+            model="claude-2.1",
             messages=[
                 {"content": "respond only in Yoda speak.", "role": "system"},
                 {"content": "what is 2 + 2?", "role": "user"},
@@ -53,7 +59,7 @@ def anthropic(streaming=False):
 # @with_langtrace_root_span("Litellm Example OpenAI Async Streaming")
 async def async_anthropic(streaming=False):
     response = await acompletion(
-        model="claude-2",
+        model="claude-2.1",
         messages=[{"content": "Hello, how are you?", "role": "user"}],
         stream=streaming,
         stream_options={"include_usage": True},
@@ -68,7 +74,25 @@ async def async_anthropic(streaming=False):
         return response
 
 
+def cohere(streaming=False):
+    response = completion(
+        model="command-r",
+        messages=[
+            {"content": "respond only in Yoda speak.", "role": "system"},
+            {"content": "Hello, how are you?", "role": "user"},
+        ],
+        stream=streaming,
+        stream_options={"include_usage": True},
+    )
+    if streaming:
+        for _ in response:
+            pass
+    else:
+        return response
+
+
 if __name__ == "__main__":
     # openAI()
-    # anthropic(streaming=False)
-    asyncio.run(async_anthropic(streaming=True))
+    anthropic(streaming=False)
+    cohere(streaming=True)
+    # asyncio.run(async_anthropic(streaming=True))