Skip to content

fix: add pytidb[models] deps to memory example#101

Merged
Mini256 merged 2 commits intopingcap:mainfrom
yahonda:add_litellm_to_memory
Jun 16, 2025
Merged

fix: add pytidb[models] deps to memory example#101
Mini256 merged 2 commits intopingcap:mainfrom
yahonda:add_litellm_to_memory

Conversation

@yahonda
Copy link
Copy Markdown
Member

@yahonda yahonda commented Jun 13, 2025

This commit adds litellm to memory example to address ModuleNotFoundError: No module named 'litellm' error.

Steps to reproduce

Follow the steps described at https://github.com/pingcap/pytidb/tree/main/examples/memory#how-to-run

Expected behavior

python main.py should run without errors.

Actual behavior without this fix

It raises the "ModuleNotFoundError: No module named 'litellm'"

(.venv) ~/src/github.com/pingcap/pytidb/examples/memory % python main.py

Traceback (most recent call last):
  File "/Users/yahonda@pingcap.com/src/github.com/pingcap/pytidb/examples/memory/main.py", line 22, in <module>
    embedding_fn = EmbeddingFunction(
        model_name="text-embedding-3-small", api_key=os.getenv("OPENAI_API_KEY")
    )
  File "/Users/yahonda@pingcap.com/src/github.com/pingcap/pytidb/examples/memory/.venv/lib/python3.13/site-packages/pytidb/embeddings/litellm.py", line 69, in __init__
    self.dimensions = len(self.get_query_embedding("test"))
                          ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^
  File "/Users/yahonda@pingcap.com/src/github.com/pingcap/pytidb/examples/memory/.venv/lib/python3.13/site-packages/pytidb/embeddings/litellm.py", line 72, in get_query_embedding
    embeddings = get_embeddings(
        api_key=self.api_key,
    ...<4 lines>...
        input=[query],
    )
  File "/Users/yahonda@pingcap.com/src/github.com/pingcap/pytidb/examples/memory/.venv/lib/python3.13/site-packages/pytidb/embeddings/litellm.py", line 29, in get_embeddings
    from litellm import embedding
ModuleNotFoundError: No module named 'litellm'
(.venv) ~/src/github.com/pingcap/pytidb/examples/memory %

Actual behavior with this fix

It works as expected that Chat with AI shows the prompt below.

(.venv) ~/src/github.com/pingcap/pytidb/examples/memory % python main.py

Chat with AI (type 'exit' to quit)
You:

@yahonda yahonda marked this pull request as ready for review June 13, 2025 05:41
@Mini256
Copy link
Copy Markdown
Member

Mini256 commented Jun 16, 2025

Using pytidb[models] instead.

@Mini256 Mini256 merged commit f720e19 into pingcap:main Jun 16, 2025
1 of 2 checks passed
@Mini256 Mini256 changed the title fix: Add litellm to memory example fix: add pytidb[models] deps to memory example Jun 16, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants