File size: 1,023 Bytes
329e3b3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
import os
import cohere.core
import instructor
import cohere
import instructor


from pydantic import BaseModel, Field

from typing import List, Dict
from pydantic import BaseModel


class Scene(BaseModel):
    narration: str
    image_prompts: List[str]


class VideoOutput(BaseModel):
    scenes: List[Scene]


# Patching the Cohere client with the instructor for enhanced capabilities
client = instructor.from_cohere(
    cohere.Client(os.environ.get("COHERE_API", "RANDOM_STRING")),
    max_tokens=5000,
    model="command-r-plus",
)


# Now, we can use the response_model parameter using only a base model
# rather than having to use the OpenAISchema class


def chatbot(prompt: str, model: str = "command-r-plus"):

    response: VideoOutput = client.chat.completions.create(
        model=model,
        max_tokens=5000,
        response_model=VideoOutput,
        messages=[
            {
                "role": "user",
                "content": prompt,
            },
        ],
    )
    return response.dict()