Repository: run-llama/llama_index
Pull request: Add support for prompt kwarg
# Description

Currently, the `llm.complete` function requires that a prompt be provided as the first positional argument (e.g., `llm.complete("my prompt)`). It does not support passing the prompt as a keyword argument (e.g., `llm.complete(prompt="my prompt")`), and throws an exception. In this PR, I've added a function that will extract the prompt in either case. Feel free to weigh in if you think this change necessitates a documentation update anywhere.

Fixes #13056, #13404

## New Package?

Did I fill in the `tool.llamahub` section in the `pyproject.toml` and provide a detailed README.md for my new integration or package?

- [ ] Yes
- [x] No

## Version Bump?

Did I bump the version in the `pyproject.toml` file of the package I am updating? (Except for the `llama-index-core` package)

- [ ] Yes
- [x] No

## Type of Change

Please delete options that are not relevant.

- [x] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

- [x] Added new unit/integration tests
- [ ] Added new notebook (that tests end-to-end)
- [x] I stared at the code and made sure it makes sense

## Suggested Checklist:

- [x] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] I have added Google Colab support for the newly added notebooks.
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Bug]: tuple index out of range when calling LLM with prompt keyword arg
### Bug Description

This succeeds:
```python
output = llm.complete(prompt)
```
but see how it fails in Steps to Reproduce below. I could be wrong if providing `prompt` as a keyword arg is unsupported, but I didn't see any documentation or comments about that.

### Version

0.10.30

### Steps to Reproduce

```python
output = llm.complete(prompt=prompt)
```


### Relevant Logs/Tracbacks

```shell
File "…/main.py", line 411, in …
    output = llm.complete(prompt=prompt)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "…/.venv/lib/python3.12/site-packages/llama_index/core/llms/callbacks.py", line 305, in wrapped_llm_predict
    prompt = kwargs.get("prompt", default=args[0])
                                          ~~~~^^^
IndexError: tuple index out of range
```

Hints from issue discussion:
<!-- Greeting -->
Hey there, @erictse! 👋 I'm [Dosu](https://dosu.dev), a bot crafted to assist with squashing bugs, answering queries, and guiding you towards becoming a contributor in the LlamaIndex universe. While we await the arrival of a human maintainer, I'm here to lend a robotic hand. I've caught your issue and am currently analyzing the circuits to provide you with a detailed response shortly. Stay tuned!



Yea it's technically not supported/advised. It might not be mentioned in the docs, but also, I don't think there's any example that uses it as a kwarg either 😅 
<!-- Answer -->
It appears that the `llm.complete` method in the LlamaIndex library version 0.10.30 expects the `prompt` to be passed as a positional argument, not as a keyword argument. This is why you're encountering the "tuple index out of range" error when attempting to use `prompt=prompt`. The method's implementation and the documentation do not support `prompt` as a keyword argument, which aligns with the maintainer's feedback that it's not advised or supported.

To avoid this error, you should call the `complete` method by directly passing the `prompt` as a positional argument, like so:

```python
output = llm.complete(prompt)
```

If you need to pass additional options, they should be included as keyword arguments after the `prompt`:

```python
output = llm.complete(prompt, formatted=False, additional_option="value")
```

This approach ensures compatibility with the current method signature and avoids the error you've encountered.


<!-- Dosu Comment Footer -->
<details>
<summary>

#### Details
</summary>

**Sources**
- [llama-index-integrations/llms/llama-index-llms-predibase/llama_index/llms/predibase/base.py](https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/llms/llama-index-llms-predibase/llama_index/llms/predibase/base.py)
- [llama-index-integrations/llms/llama-index-llms-bedrock/llama_index/llms/bedrock/base.py](https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/llms/llama-index-llms-bedrock/llama_index/llms/bedrock/base.py)
- [llama-index-core/llama_index/core/llms/llm.py](https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/llms/llm.py)
- [llama-index-legacy/llama_index/legacy/llms/llm.py](https://github.com/run-llama/llama_index/blob/main/llama-index-legacy/llama_index/legacy/llms/llm.py)
- [llama-index-core/llama_index/core/llms/callbacks.py](https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/llms/callbacks.py)

<br/>

**About Dosu**
This response is meant to be useful and save you time. It is **not meant to be a precise solution**, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.
</details>

*To continue the conversation, mention @dosu.*

Commit message: Add support for prompt kwarg

Patch to review:
diff --git a/llama-index-core/llama_index/core/llms/callbacks.py b/llama-index-core/llama_index/core/llms/callbacks.py
index 76ebf7bc9a..048715d611 100644
--- a/llama-index-core/llama_index/core/llms/callbacks.py
+++ b/llama-index-core/llama_index/core/llms/callbacks.py
@@ -236,9 +236,18 @@ def wrapper_logic(_self: Any) -> Generator[CallbackManager, None, None]:
 
             yield callback_manager
 
+        def extract_prompt(*args: Any, **kwargs: Any) -> str:
+            if len(args) > 0:
+                return str(args[0])
+            elif "prompt" in kwargs:
+                return kwargs["prompt"]
+            else:
+                return str(kwargs)
+
         async def wrapped_async_llm_predict(
             _self: Any, *args: Any, **kwargs: Any
         ) -> Any:
+            prompt = extract_prompt(*args, **kwargs)
             with wrapper_logic(_self) as callback_manager:
                 span_id = dispatcher.root.current_span_id or ""
                 model_dict = _self.to_dict()
@@ -246,7 +255,7 @@ async def wrapped_async_llm_predict(
                 dispatcher.event(
                     LLMCompletionStartEvent(
                         model_dict=model_dict,
-                        prompt=str(args[0]),
+                        prompt=prompt,
                         additional_kwargs=kwargs,
                         span_id=span_id,
                     )
@@ -254,7 +263,7 @@ async def wrapped_async_llm_predict(
                 event_id = callback_manager.on_event_start(
                     CBEventType.LLM,
                     payload={
-                        EventPayload.PROMPT: args[0],
+                        EventPayload.PROMPT: prompt,
                         EventPayload.ADDITIONAL_KWARGS: kwargs,
                         EventPayload.SERIALIZED: _self.to_dict(),
                     },
@@ -269,7 +278,7 @@ async def wrapped_gen() -> CompletionResponseAsyncGen:
                         async for x in f_return_val:
                             dispatcher.event(
                                 LLMCompletionEndEvent(
-                                    prompt=str(args[0]),
+                                    prompt=prompt,
                                     response=x,
                                     span_id=span_id,
                                 )
@@ -280,7 +289,7 @@ async def wrapped_gen() -> CompletionResponseAsyncGen:
                         callback_manager.on_event_end(
                             CBEventType.LLM,
                             payload={
-                                EventPayload.PROMPT: args[0],
+                                EventPayload.PROMPT: prompt,
                                 EventPayload.COMPLETION: last_response,
                             },
                             event_id=event_id,
@@ -291,14 +300,14 @@ async def wrapped_gen() -> CompletionResponseAsyncGen:
                     callback_manager.on_event_end(
                         CBEventType.LLM,
                         payload={
-                            EventPayload.PROMPT: args[0],
+                            EventPayload.PROMPT: prompt,
                             EventPayload.COMPLETION: f_return_val,
                         },
                         event_id=event_id,
                     )
                     dispatcher.event(
                         LLMCompletionEndEvent(
-                            prompt=str(args[0]),
+                            prompt=prompt,
                             response=f_return_val,
                             span_id=span_id,
                         )
@@ -307,6 +316,7 @@ async def wrapped_gen() -> CompletionResponseAsyncGen:
             return f_return_val
 
         def wrapped_llm_predict(_self: Any, *args: Any, **kwargs: Any) -> Any:
+            prompt = extract_prompt(*args, **kwargs)
             with wrapper_logic(_self) as callback_manager:
                 span_id = dispatcher.root.current_span_id or ""
                 model_dict = _self.to_dict()
@@ -314,7 +324,7 @@ def wrapped_llm_predict(_self: Any, *args: Any, **kwargs: Any) -> Any:
                 dispatcher.event(
                     LLMCompletionStartEvent(
                         model_dict=model_dict,
-                        prompt=str(args[0]),
+                        prompt=prompt,
                         additional_kwargs=kwargs,
                         span_id=span_id,
                     )
@@ -322,7 +332,7 @@ def wrapped_llm_predict(_self: Any, *args: Any, **kwargs: Any) -> Any:
                 event_id = callback_manager.on_event_start(
                     CBEventType.LLM,
                     payload={
-                        EventPayload.PROMPT: args[0],
+                        EventPayload.PROMPT: prompt,
                         EventPayload.ADDITIONAL_KWARGS: kwargs,
                         EventPayload.SERIALIZED: _self.to_dict(),
                     },
@@ -336,7 +346,7 @@ def wrapped_gen() -> CompletionResponseGen:
                         for x in f_return_val:
                             dispatcher.event(
                                 LLMCompletionEndEvent(
-                                    prompt=str(args[0]), response=x, span_id=span_id
+                                    prompt=prompt, response=x, span_id=span_id
                                 )
                             )
                             yield cast(CompletionResponse, x)
@@ -345,7 +355,7 @@ def wrapped_gen() -> CompletionResponseGen:
                         callback_manager.on_event_end(
                             CBEventType.LLM,
                             payload={
-                                EventPayload.PROMPT: args[0],
+                                EventPayload.PROMPT: prompt,
                                 EventPayload.COMPLETION: last_response,
                             },
                             event_id=event_id,
@@ -356,14 +366,14 @@ def wrapped_gen() -> CompletionResponseGen:
                     callback_manager.on_event_end(
                         CBEventType.LLM,
                         payload={
-                            EventPayload.PROMPT: args[0],
+                            EventPayload.PROMPT: prompt,
                             EventPayload.COMPLETION: f_return_val,
                         },
                         event_id=event_id,
                     )
                     dispatcher.event(
                         LLMCompletionEndEvent(
-                            prompt=str(args[0]),
+                            prompt=prompt,
                             response=f_return_val,
                             span_id=span_id,
                         )
diff --git a/llama-index-core/tests/llms/test_callbacks.py b/llama-index-core/tests/llms/test_callbacks.py
new file mode 100644
index 0000000000..42177bbe4b
--- /dev/null
+++ b/llama-index-core/tests/llms/test_callbacks.py
@@ -0,0 +1,57 @@
+import pytest
+from llama_index.core.llms.llm import LLM
+from llama_index.core.llms.mock import MockLLM
+
+
+@pytest.fixture()
+def llm() -> LLM:
+    return MockLLM()
+
+
+@pytest.fixture()
+def prompt() -> str:
+    return "test prompt"
+
+
+def test_llm_complete_prompt_arg(llm: LLM, prompt: str) -> None:
+    res = llm.complete(prompt)
+    expected_res_text = prompt
+    assert res.text == expected_res_text
+
+
+def test_llm_complete_prompt_kwarg(llm: LLM, prompt: str) -> None:
+    res = llm.complete(prompt=prompt)
+    expected_res_text = prompt
+    assert res.text == expected_res_text
+
+
+def test_llm_complete_throws_if_duplicate_prompt(llm: LLM, prompt: str) -> None:
+    with pytest.raises(TypeError):
+        llm.complete(prompt, prompt=prompt)
+
+
+def test_llm_complete_throws_if_no_prompt(llm: LLM) -> None:
+    with pytest.raises(TypeError):
+        llm.complete()
+
+
+def test_llm_stream_complete_prompt_arg(llm: LLM, prompt: str) -> None:
+    res_text = "".join(r.delta for r in llm.stream_complete(prompt))
+    expected_res_text = prompt
+    assert res_text == expected_res_text
+
+
+def test_llm_stream_complete_prompt_kwarg(llm: LLM, prompt: str) -> None:
+    res_text = "".join(r.delta for r in llm.stream_complete(prompt=prompt))
+    expected_res_text = prompt
+    assert res_text == expected_res_text
+
+
+def test_llm_stream_complete_throws_if_duplicate_prompt(llm: LLM, prompt: str) -> None:
+    with pytest.raises(TypeError):
+        llm.stream_complete(prompt, prompt=prompt)
+
+
+def test_llm_stream_complete_throws_if_no_prompt(llm: LLM) -> None:
+    with pytest.raises(TypeError):
+        llm.stream_complete()


Code review comments:
diff_hunk: @@ -236,25 +236,34 @@ def wrapper_logic(_self: Any) -> Generator[CallbackManager, None, None]:
 
             yield callback_manager
 
+        def extract_prompt(*args: Any, **kwargs: Any) -> str:
+            if len(args) > 0:
+                return str(args[0])
+            elif "prompt" in kwargs:
+                return kwargs["prompt"]
+            else:
+                return str(kwargs)
path: llama-index-core/llama_index/core/llms/callbacks.py
text: @user1:
Why we want to cast `kwargs` to str? 🤔 

@author:
Logan suggested this, I guess it is just one way to avoid an exception. The obvious alternative is raising a ValueError or some other exception. I think that might be more appropriate, since this else block covers unexpected / improper usage, but let me know what you think!

@user1:
I would vote for ValueError 🙋  
```suggestion
raise ValueError("No prompt provided in positional or keyword arguments")
```

@user2:
I'm not sure this would even be possible to hit (python should catch this earlier) -- but either way works for me