rags

1
"""Streaming callback manager."""
2
from llama_index.callbacks.base_handler import BaseCallbackHandler
3
from llama_index.callbacks.schema import CBEventType
4

5
from typing import Optional, Dict, Any, List, Callable
6

7
STORAGE_DIR = "./storage"  # directory to cache the generated index
8
DATA_DIR = "./data"  # directory containing the documents to index
9

10

11
class StreamlitFunctionsCallbackHandler(BaseCallbackHandler):
12
    """Callback handler that outputs streamlit components given events."""
13

14
    def __init__(self, msg_handler: Callable[[str], Any]) -> None:
15
        """Initialize the base callback handler."""
16
        self.msg_handler = msg_handler
17
        super().__init__([], [])
18

19
    def on_event_start(
20
        self,
21
        event_type: CBEventType,
22
        payload: Optional[Dict[str, Any]] = None,
23
        event_id: str = "",
24
        parent_id: str = "",
25
        **kwargs: Any,
26
    ) -> str:
27
        """Run when an event starts and return id of event."""
28
        if event_type == CBEventType.FUNCTION_CALL:
29
            if payload is None:
30
                raise ValueError("Payload cannot be None")
31
            arguments_str = payload["function_call"]
32
            tool_str = payload["tool"].name
33
            print_str = f"Calling function: {tool_str} with args: {arguments_str}\n\n"
34
            self.msg_handler(print_str)
35
        else:
36
            pass
37
        return event_id
38

39
    def on_event_end(
40
        self,
41
        event_type: CBEventType,
42
        payload: Optional[Dict[str, Any]] = None,
43
        event_id: str = "",
44
        **kwargs: Any,
45
    ) -> None:
46
        """Run when an event ends."""
47
        pass
48
        # TODO: currently we don't need to do anything here
49
        # if event_type == CBEventType.FUNCTION_CALL:
50
        #     response = payload["function_call_response"]
51
        #     # Add this to queue
52
        #     print_str = (
53
        #         f"\n\nGot output: {response}\n"
54
        #         "========================\n\n"
55
        #     )
56
        # elif event_type == CBEventType.AGENT_STEP:
57
        #     # put response into queue
58
        #     self._queue.put(payload["response"])
59

60
    def start_trace(self, trace_id: Optional[str] = None) -> None:
61
        """Run when an overall trace is launched."""
62
        pass
63

64
    def end_trace(
65
        self,
66
        trace_id: Optional[str] = None,
67
        trace_map: Optional[Dict[str, List[str]]] = None,
68
    ) -> None:
69
        """Run when an overall trace is exited."""
70
        pass
71