Skip to content

Commit

Permalink
Jerry/fix llms callbacks issue (#7165)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jerryjliu
jerryjliu committed Aug 6, 2023
1 parent f089096 commit 8629e50
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 46 deletions.
39 changes: 3 additions & 36 deletions docs/examples/query_engine/RouterQueryEngine.ipynb
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
Expand All @@ -10,7 +9,6 @@
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
Expand Down Expand Up @@ -64,7 +62,6 @@
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
Expand Down Expand Up @@ -106,7 +103,6 @@
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
Expand Down Expand Up @@ -135,7 +131,6 @@
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
Expand Down Expand Up @@ -176,7 +171,6 @@
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
Expand Down Expand Up @@ -253,41 +247,15 @@
]
},
{
"cell_type": "code",
"execution_count": 12,
"cell_type": "markdown",
"metadata": {
"scrolled": false,
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Selecting query engine 1: The question is asking for specific context from Paul Graham's essay on What I Worked On..\n",
"> [retrieve] Total LLM token usage: 0 tokens\n",
"> [retrieve] Total embedding token usage: 9 tokens\n",
"> [get_response] Total LLM token usage: 1990 tokens\n",
"> [get_response] Total embedding token usage: 0 tokens\n"
]
},
{
"data": {
"text/plain": [
"Response(response=\"\\nAfter RICS, Paul Graham decided to focus on his mother's health and well-being. He flew up to Oregon to visit her regularly and used the time to think about his future. He asked Jessica if she wanted to be president of YC, but when she declined, he decided to recruit Sam Altman. He spent the rest of 2013 leaving running YC more and more to Sam, and focusing on his mother's cancer. After his mother passed away in January 2014, Paul Graham decided to take a break from YC and focus on painting. He spent most of 2014 painting and writing essays, and in March 2015 he started working on Lisp again.\", source_nodes=[NodeWithScore(node=Node(text=\"of my attention. It had already eaten Arc, and was in the process of eating essays too. Either YC was my life's work or I'd have to leave eventually. And it wasn't, so I would.\\n\\nIn the summer of 2012 my mother had a stroke, and the cause turned out to be a blood clot caused by colon cancer. The stroke destroyed her balance, and she was put in a nursing home, but she really wanted to get out of it and back to her house, and my sister and I were determined to help her do it. I used to fly up to Oregon to visit her regularly, and I had a lot of time to think on those flights. On one of them I realized I was ready to hand YC over to someone else.\\n\\nI asked Jessica if she wanted to be president, but she didn't, so we decided we'd try to recruit Sam Altman. We talked to Robert and Trevor and we agreed to make it a complete changing of the guard. Up till that point YC had been controlled by the original LLC we four had started. But we wanted YC to last for a long time, and to do that it couldn't be controlled by the founders. So if Sam said yes, we'd let him reorganize YC. Robert and I would retire, and Jessica and Trevor would become ordinary partners.\\n\\nWhen we asked Sam if he wanted to be president of YC, initially he said no. He wanted to start a startup to make nuclear reactors. But I kept at it, and in October 2013 he finally agreed. We decided he'd take over starting with the winter 2014 batch. For the rest of 2013 I left running YC more and more to Sam, partly so he could learn the job, and partly because I was focused on my mother, whose cancer had returned.\\n\\nShe died on January 15, 2014. We knew this was coming, but it was still hard when it did.\\n\\nI kept working on YC till March, to help get that batch of startups through Demo Day, then I checked out pretty completely. (I still talk to alumni and to new startups working on things I'm interested in, but that only takes a few hours a week.)\\n\\nWhat should I do next? Rtm's advice hadn't included anything about that. I wanted to do something completely different, so I decided I'd paint. I wanted to see how good I could get if I really focused on it. So the day after I stopped working on YC, I started painting. I was rusty and it took a while to get back into shape, but it was at least completely engaging. [18]\\n\\nI spent most of the rest of 2014 painting. I'd never been able to work so uninterruptedly before, and I got to be better than I had been. Not good enough, but better. Then in November, right in the middle of a painting, I ran out of steam. Up till that point I'd always been curious to see how the painting I was working on would turn out, but suddenly finishing this one seemed like a chore. So I stopped working on it and cleaned my brushes and haven't painted since. So far anyway.\\n\\nI realize that sounds rather wimpy. But attention is a zero sum game. If you can choose what to work on, and you choose a project that's not the best one (or at least a good one) for you, then it's getting in the way of another project that is. And at 50 there was some opportunity cost to screwing around.\\n\\nI started writing essays again, and wrote a bunch of new ones over the next few months. I even wrote a couple that weren't about startups. Then in March 2015 I started working on Lisp again.\\n\\nThe distinctive thing about Lisp is that its core is a language defined by writing an interpreter in itself. It wasn't originally intended as a programming language in the ordinary sense. It was meant to be a formal model of computation, an alternative to the Turing machine. If you want to write an interpreter for a language in itself, what's the minimum set of predefined operators you need? The Lisp that John McCarthy invented, or more accurately discovered, is an answer to that question. [19]\\n\\nMcCarthy didn't realize this Lisp could even be used to program computers\", doc_id='cfd50785-c54a-4e07-b474-561541968da9', embedding=None, doc_hash='b6524989f50c19316fcc4135d476deedc62b79ab141d9f650743a6fe5f3558c9', extra_info=None, node_info={'start': 60170, 'end': 64027, '_node_type': <NodeType.TEXT: '1'>}, relationships={<DocumentRelationship.SOURCE: '1'>: 'af3dc326-313b-40f5-87ad-007838e7a370', <DocumentRelationship.PREVIOUS: '2'>: 'ab7ed037-4269-4593-b5ff-0ce3d9213cbd', <DocumentRelationship.NEXT: '3'>: '14d14357-e5cf-4015-8a65-2cc9fd3fb5c4'}), score=0.7906612296986986), NodeWithScore(node=Node(text='of customers almost entirely from among their batchmates.\\n\\nI had not originally intended YC to be a full-time job. I was going to do three things: hack, write essays, and work on YC. As YC grew, and I grew more excited about it, it started to take up a lot more than a third of my attention. But for the first few years I was still able to work on other things.\\n\\nIn the summer of 2006, Robert and I started working on a new version of Arc. This one was reasonably fast, because it was compiled into Scheme. To test this new Arc, I wrote Hacker News in it. It was originally meant to be a news aggregator for startup founders and was called Startup News, but after a few months I got tired of reading about nothing but startups. Plus it wasn\\'t startup founders we wanted to reach. It was future startup founders. So I changed the name to Hacker News and the topic to whatever engaged one\\'s intellectual curiosity.\\n\\nHN was no doubt good for YC, but it was also by far the biggest source of stress for me. If all I\\'d had to do was select and help founders, life would have been so easy. And that implies that HN was a mistake. Surely the biggest source of stress in one\\'s work should at least be something close to the core of the work. Whereas I was like someone who was in pain while running a marathon not from the exertion of running, but because I had a blister from an ill-fitting shoe. When I was dealing with some urgent problem during YC, there was about a 60% chance it had to do with HN, and a 40% chance it had do with everything else combined. [17]\\n\\nAs well as HN, I wrote all of YC\\'s internal software in Arc. But while I continued to work a good deal in Arc, I gradually stopped working on Arc, partly because I didn\\'t have time to, and partly because it was a lot less attractive to mess around with the language now that we had all this infrastructure depending on it. So now my three projects were reduced to two: writing essays and working on YC.\\n\\nYC was different from other kinds of work I\\'ve done. Instead of deciding for myself what to work on, the problems came to me. Every 6 months there was a new batch of startups, and their problems, whatever they were, became our problems. It was very engaging work, because their problems were quite varied, and the good founders were very effective. If you were trying to learn the most you could about startups in the shortest possible time, you couldn\\'t have picked a better way to do it.\\n\\nThere were parts of the job I didn\\'t like. Disputes between cofounders, figuring out when people were lying to us, fighting with people who maltreated the startups, and so on. But I worked hard even at the parts I didn\\'t like. I was haunted by something Kevin Hale once said about companies: \"No one works harder than the boss.\" He meant it both descriptively and prescriptively, and it was the second part that scared me. I wanted YC to be good, so if how hard I worked set the upper bound on how hard everyone else worked, I\\'d better work very hard.\\n\\nOne day in 2010, when he was visiting California for interviews, Robert Morris did something astonishing: he offered me unsolicited advice. I can only remember him doing that once before. One day at Viaweb, when I was bent over double from a kidney stone, he suggested that it would be a good idea for him to take me to the hospital. That was what it took for Rtm to offer unsolicited advice. So I remember his exact words very clearly. \"You know,\" he said, \"you should make sure Y Combinator isn\\'t the last cool thing you do.\"\\n\\nAt the time I didn\\'t understand what he meant, but gradually it dawned on me that he was saying I should quit. This seemed strange advice, because YC was doing great. But if there was one thing rarer than Rtm offering advice, it was Rtm being wrong. So this set me thinking. It was true that on my current trajectory, YC would be the last thing I did, because it was only taking up more of my attention. It had already eaten Arc, and was in the process of eating', doc_id='ab7ed037-4269-4593-b5ff-0ce3d9213cbd', embedding=None, doc_hash='28da476588fa6a7c04e3fc8d0c8490de5a6aa3f4b46ada11723bd524402b1d33', extra_info=None, node_info={'start': 56162, 'end': 60161, '_node_type': <NodeType.TEXT: '1'>}, relationships={<DocumentRelationship.SOURCE: '1'>: 'af3dc326-313b-40f5-87ad-007838e7a370', <DocumentRelationship.PREVIOUS: '2'>: 'f9e74e9d-cdd9-43c4-8742-76c38200305f', <DocumentRelationship.NEXT: '3'>: 'cfd50785-c54a-4e07-b474-561541968da9'}), score=0.7863717426739052)], extra_info={'cfd50785-c54a-4e07-b474-561541968da9': None, 'ab7ed037-4269-4593-b5ff-0ce3d9213cbd': None})"
]
},
"execution_count": 12,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"query_engine.query(\"What did Paul Graham do after RICS?\")"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
Expand Down Expand Up @@ -374,7 +342,6 @@
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
Expand Down Expand Up @@ -479,7 +446,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
Expand All @@ -493,7 +460,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.6"
"version": "3.10.10"
}
},
"nbformat": 4,
Expand Down
24 changes: 14 additions & 10 deletions llama_index/llms/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,13 +91,15 @@ def wrapper_logic(_self: Any) -> Generator[CallbackManager, None, None]:

yield callback_manager

async def wrapped_async_llm_chat(_self: Any, *args: Any, **kwargs: Any) -> Any:
async def wrapped_async_llm_chat(
_self: Any, messages: Sequence[ChatMessage], **kwargs: Any
) -> Any:
with wrapper_logic(_self) as callback_manager:
event_id = callback_manager.on_event_start(
CBEventType.LLM, payload={EventPayload.MESSAGES: args[0]}
CBEventType.LLM, payload={EventPayload.MESSAGES: messages}
)

f_return_val = await f(_self, *args, **kwargs)
f_return_val = await f(_self, messages, **kwargs)
if isinstance(f_return_val, AsyncGenerator):
# intercept the generator and add a callback to the end
async def wrapped_gen() -> ChatResponseAsyncGen:
Expand All @@ -109,7 +111,7 @@ async def wrapped_gen() -> ChatResponseAsyncGen:
callback_manager.on_event_end(
CBEventType.LLM,
payload={
EventPayload.MESSAGES: args[0],
EventPayload.MESSAGES: messages,
EventPayload.RESPONSE: last_response,
},
event_id=event_id,
Expand All @@ -120,20 +122,22 @@ async def wrapped_gen() -> ChatResponseAsyncGen:
callback_manager.on_event_end(
CBEventType.LLM,
payload={
EventPayload.MESSAGES: args[0],
EventPayload.MESSAGES: messages,
EventPayload.RESPONSE: f_return_val,
},
event_id=event_id,
)

return f_return_val

def wrapped_llm_chat(_self: Any, *args: Any, **kwargs: Any) -> Any:
def wrapped_llm_chat(
_self: Any, messages: Sequence[ChatMessage], **kwargs: Any
) -> Any:
with wrapper_logic(_self) as callback_manager:
event_id = callback_manager.on_event_start(
CBEventType.LLM, payload={EventPayload.MESSAGES: args[0]}
CBEventType.LLM, payload={EventPayload.MESSAGES: messages}
)
f_return_val = f(_self, *args, **kwargs)
f_return_val = f(_self, messages, **kwargs)

if isinstance(f_return_val, Generator):
# intercept the generator and add a callback to the end
Expand All @@ -146,7 +150,7 @@ def wrapped_gen() -> ChatResponseGen:
callback_manager.on_event_end(
CBEventType.LLM,
payload={
EventPayload.MESSAGES: args[0],
EventPayload.MESSAGES: messages,
EventPayload.RESPONSE: last_response,
},
event_id=event_id,
Expand All @@ -157,7 +161,7 @@ def wrapped_gen() -> ChatResponseGen:
callback_manager.on_event_end(
CBEventType.LLM,
payload={
EventPayload.MESSAGES: args[0],
EventPayload.MESSAGES: messages,
EventPayload.RESPONSE: f_return_val,
},
event_id=event_id,
Expand Down

0 comments on commit 8629e50

Please sign in to comment.