wehub-resource-sync
38 lines
860 B
Python
38 lines
860 B
Python
from llama_index.core.workflow import (
|
|
Event,
|
|
StartEvent,
|
|
StopEvent,
|
|
Workflow,
|
|
step,
|
|
)
|
|
from llama_index.llms.openai import OpenAI
|
|
|
|
import mlflow
|
|
|
|
|
|
class JokeEvent(Event):
|
|
joke: str
|
|
|
|
|
|
class JokeFlow(Workflow):
|
|
llm = OpenAI()
|
|
|
|
@step
|
|
async def generate_joke(self, ev: StartEvent) -> JokeEvent:
|
|
topic = ev.topic
|
|
prompt = f"Write your best joke about {topic}."
|
|
response = await self.llm.acomplete(prompt)
|
|
return JokeEvent(joke=str(response))
|
|
|
|
@step
|
|
async def critique_joke(self, ev: JokeEvent) -> StopEvent:
|
|
joke = ev.joke
|
|
|
|
prompt = f"Give a thorough analysis and critique of the following joke: {joke}"
|
|
response = await self.llm.acomplete(prompt)
|
|
return StopEvent(result=str(response))
|
|
|
|
|
|
w = JokeFlow(timeout=10, verbose=False)
|
|
mlflow.models.set_model(w)
|