AidenS33/llamaHack

1import os2from fastapi import FastAPI3from pydantic import BaseModel4from llama_index.indices.managed.llama_cloud import LlamaCloudIndex5from llama_index.llms.cohere import Cohere6 7app = FastAPI()8 9class PokemonBattleRequest(BaseModel):10    user_pokemon: str11    user_pokemon_type: str12    opponent_pokemon: str13    opponent_pokemon_type: str14 15cohere_api_key = os.getenv("COHERE_API_KEY")16cohere_llm = Cohere(api_key=cohere_api_key)17llama_cloud_api_key = os.getenv("LLAMA_CLOUD_API_KEY")18 19type_index = LlamaCloudIndex(20    name="pokemonTypes1Test",21    project_name="Default",22    organization_id="b33f284d-02d9-4a32-97ca-7031a4881df2",23    api_key=llama_cloud_api_key24)25 26move_index = LlamaCloudIndex(27    name="pokemonmoveswithtypes",28    project_name="Default",29    organization_id="b33f284d-02d9-4a32-97ca-7031a4881df2",30    api_key=llama_cloud_api_key31)32 33# Function to get strong typing34def get_strong_typing(user_pokemon_type, opponent_pokemon_type):35    query = f"What type is strong against {opponent_pokemon_type}-type Pokémon?"36    response = type_index.as_query_engine(llm=cohere_llm).query(query)37    if response and hasattr(response, "response"):38        return response.response39    else:40        return None41 42# Function to get recommended move43def get_recommended_move(user_pokemon, strong_type):44    query = f"What move should {user_pokemon} use that is {strong_type}-type?"45    response = move_index.as_query_engine(llm=cohere_llm).query(query)46    if response and hasattr(response, "response"):47        return response.response48    else:49        return None50 51# API endpoint52@app.post("/suggest_move")53async def suggest_move(request: PokemonBattleRequest):54    strong_type = get_strong_typing(request.user_pokemon_type.lower(), request.opponent_pokemon_type.lower())55    if strong_type:56        recommended_move = get_recommended_move(request.user_pokemon.lower(), strong_type)57        if recommended_move:58            return {59                "user_pokemon": request.user_pokemon,60                "recommended_move": recommended_move,61                "strong_type": strong_type62            }63        else:64            return {"error": "No recommended move found."}65    else:66        return {"error": "No strong typing found."}67