from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
# pip install llama-index-indices-managed-llama-cloud

index = LlamaCloudIndex(
  name="My Index", 
  project_name="My Project",
  organization_id="e793a802-cb91-4e6a-bd49-61d0ba2ac5f9",
  api_key="llx-..."
)

# retrieve nodes
nodes = index.as_retriever().retrieve(query)
# or use as a query engine
response = index.as_query_engine().query(query)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

documents = SimpleDirectoryReader("data").load_data()

index = VectorStoreIndex.from_documents(documents)

query_engine = index.as_query_engine()

response = query_engine.query("What did the author do growing up?")

print(response)

# set up list of tools
tools = [
    QueryEngineTool(
        query_engine=vector_query_engine,
        metadata=ToolMetadata(
            # this is how the LLM knows what the tool does
            name="pg_essay",
            description="Paul Graham essay on What I Worked On",
        ),
    ),
    # more query engines or other tools go here
]

# define sample Tool
def multiply(a: int, b: int) -> int:
    """Multiply two integers and returns the result integer"""
    return a * b

multiply_tool = FunctionTool.from_defaults(fn=multiply)

# initialize ReAct agent
agent = ReActAgent.from_tools(
  [
    multiply_tool
    # other tools here
  ], 
  verbose=True
)

from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec

tool_spec = DuckDuckGoSearchToolSpec()

# initialize ReAct agent
# with a list of tools from the spec
agent = ReActAgent.from_tools(
  DuckDuckGoSearchToolSpec.to_tool_list(),
  verbose=True
)

# ReAct agent
agent = ReActAgent.from_tools(
  [multiply_tool], 
  verbose=True
)

# OR structured planner agent
worker = FunctionCallingAgentWorker.from_tools(
    [lyft_tool, uber_tool], verbose=True
)
agent = StructuredPlannerAgent(
    worker, tools=[lyft_tool, uber_tool], verbose=True
)

# OR language agent tree search
worker = LATSAgentWorker.from_tools(
    query_engine_tools,
    llm=llm,
    num_expansions=2,
    max_rollouts=3,
    verbose=True,
)
agent = worker.as_agent()

from llama_index.llms.openai import OpenAI

class OpenAIGenerator(Workflow):
    @step()
    async def generate(self, ev: StartEvent) -> StopEvent:
        query = ev.get("query")
        llm = OpenAI()
        response = await llm.acomplete(query)
        return StopEvent(result=str(response))

w = OpenAIGenerator(timeout=10, verbose=False)
result = await w.run(query="What's LlamaIndex?")
print(result)

draw_all_possible_flows()