Fireworks AI <> Weave Integration

import weave
from fireworks.client import Fireworks

# Initialize weave project
weave.init("my_awesome_project")

client = Fireworks(api_key="<FIREWORKS_API_KEY>")

@weave.op()
def get_fireworks_completion(query):
    response = client.chat.completions.create(
        model="accounts/fireworks/models/llama-v3p1-8b-instruct",
        messages=[{
            "role": "user",
            "content": query,
        }],
    )
    return response.choices[0].message.content

print(get_fireworks_completion("Say this is a test"))

import weave
from weave import Model
from fireworks.client import Fireworks

# Initialize weave project
weave.init("my_awesome_project")

client = Fireworks(api_key="<FIREWORKS_API_KEY>")

class MyModel(Model):
    model_name: str

    @weave.op()
    def infer(self, query: str) -> str:
        response = client.chat.completions.create(
            model = model_name,
            messages = [{
                "role": "user",
                "content": query,
            }],
        )
        return response.choices[0].message.content

# Initialize Model
model = MyModel(
    model_name="accounts/fireworks/models/llama-v3p1-8b-instruct"
)

print(model.infer("Say this is a test"))