LlamaIndex agent

1"""This module defines the main entry point for the Apify LlamaIndex Agent.
2
3This Agent template is intended to give example on how to use LlamaIndex Agent with Apify Actors.
4It extracts contact details from a plain text query with a URL.
5
6Feel free to modify this file to suit your specific needs.
7
8To build Apify Actors, utilize the Apify SDK toolkit, read more at the official documentation:
9https://docs.apify.com/sdk/python
10"""
11
12from __future__ import annotations
13
14import math
15from typing import TYPE_CHECKING
16
17from apify import Actor
18from llama_index.llms.openai import OpenAI
19
20from .agent import run_agent
21
22if TYPE_CHECKING:
23    from llama_index.core.chat_engine.types import AgentChatResponse
24
25
26async def main() -> None:
27    """Main entry point for the Apify LlamaIndex Agent.
28
29    This coroutine is executed using `asyncio.run()`, so it must remain an asynchronous function for proper execution.
30    Asynchronous execution is required for communication with Apify platform, and it also enhances performance in
31    the field of web scraping significantly.
32    """
33    async with Actor:
34        Actor.log.info('Starting LlamaIndex Agent')
35        count = math.ceil((Actor.get_env().get('memory_mbytes', 1024) or 1024) // 1024)
36        await Actor.charge(event_name='actor-start-gb', count=count)
37        Actor.log.info('Charged for Actor start %d GB', count)
38        try:
39            if not (actor_input := await Actor.get_input()):
40                await Actor.fail(status_message='Actor input was not provided')
41                return
42            await check_inputs(actor_input)
43            answer = await run_query(actor_input['query'], actor_input['modelName'])
44            await Actor.push_data({'query': actor_input['query'], 'answer': answer})
45            Actor.log.info('Charging for task completed')
46            await Actor.charge(event_name='task-completed', count=1)
47        except Exception as e:
48            await Actor.fail(status_message='Failed to process query', exception=e)
49
50
51async def check_inputs(actor_input: dict) -> None:
52    """Check that provided input exists.
53
54    :raises Exception: If query is not provided
55    """
56    if not actor_input.get('query'):
57        msg = 'Input `query` is not provided. Please verify that the `query` is correctly set.'
58        await Actor.fail(status_message=msg)
59
60
61async def run_query(query: str, model_name: str) -> AgentChatResponse | None:
62    """Process query with LlamaIndex Agent."""
63    llm = OpenAI(model=str(model_name), temperature=0)
64    try:
65        return await run_agent(query=query, llm=llm, verbose=True)
66    except Exception as e:
67        msg = f'Error running LlamaIndex Agent, error: {e}'
68        await Actor.fail(status_message=msg, exception=e)