from crewai.flow.flow import Flow, listen, start
from litellm import completion

class ExampleFlow(Flow):
    model = "ollama/deepseek-r1:8b"

    @start()
    def generate_city(self):
        print("Starting flow")
        print(f"Flow State ID: {self.state['id']}")

        response = completion(
            model=self.model,
            messages=[
                {
                    "role": "user",
                    "content": "Give me the name of a random city in the world. I just need one name in the end. Your final answer should be just the name of the city. Please do not use abbreviation.",
                },
            ],
            api_base = "http://localhost:11434",
        )

        output = response["choices"][0]["message"]["content"]

        think_end = output.find('</think>')
        random_city = output[think_end+len('</think>'):].strip()

        self.state["city"] = random_city
        print(f"The output is: {output}")
        print(f"Random City: {random_city}")

    @listen(generate_city)
    def generate_fun_fact(self):
        response = completion(
            model=self.model,
            messages=[
                {
                    "role": "user",
                    "content": f"Tell me a fun fact about the city - {self.state['city']}.",
                },
            ],
            api_base = "http://localhost:11434",
        )
        output = response["choices"][0]["message"]["content"]
        think_end = output.find('</think>')
        fun_fact = output[think_end+len('</think>'):].strip()

        print(f"The output is: {output}")
        self.state["fun_fact"] = fun_fact
        return fun_fact
    
flow = ExampleFlow()
result = flow.kickoff()
print(f"Generated fun fact: {result}")