Feat/litellm llmclient abstraction

[nkanu17]

Summary

This PR replaces the custom LLM wrapper classes in llms.py with a unified LLMClient abstraction layer backed by LiteLLM. All LLM operations—chat completions, embeddings, and LangChain Embeddings instances—now go through a single entry point, reducing maintenance burden and enabling support for multiple providers without code changes.

Issue: #105

---

Problem

The previous llms.py module had several issues:

Issue	Impact
Multiple wrapper classes (OpenAIClientWrapper, AnthropicClientWrapper, BedrockClientWrapper)	Duplicated logic, maintenance overhead
Manual MODEL_NAME_MAP dictionary	Required updates for every new model
Different response parsing per provider	Inconsistent error handling
Adding a provider = new wrapper class	High friction for expansion

---

Solution

Replace custom wrappers with LLMClient, a facade backed by LiteLLM:

Before	After
3 wrapper classes	1 unified client
Manual provider detection	Automatic from model name
Custom response parsing	Standardized dataclasses
Adding provider = new class	Adding provider = just use it

Why this reduces maintenance:

1. Provider API changes are isolated - LiteLLM handles SDK updates; we just bump the dependency
2. Zero-effort provider expansion - Gemini, Cohere, Ollama work without code changes
3. Centralized bug fixes - One place for logging, retries, error handling
4. Backend swappability - LLMBackend Protocol allows replacing LiteLLM without touching application code

---

Architecture

┌─────────────────────────────────────────────────────────────────────────────┐
│                              CALL SITES                                     │
│  extraction.py | summarization.py | memory_strategies.py | api.py | main.py│
│                                                                             │
│     await LLMClient.create_chat_completion(model, messages, ...)            │
│     await LLMClient.create_embedding(model, input_texts, ...)               │
│     LLMClient.create_embeddings()  → LangChain Embeddings instance          │
│     LLMClient.get_model_config(model_name)                                  │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                            LLMClient                                        │
│                   (Facade - Public API, Never Changes)                      │
│                                                                             │
│  create_chat_completion() → ChatCompletionResponse                          │
│  create_embedding()       → EmbeddingResponse                               │
│  create_embeddings()      → LangChain Embeddings (for vector stores)        │
│  get_model_config()       → ModelConfig                                     │
│  set_backend() / reset()  → For testing                                     │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                     Internal: LiteLLM (hidden, swappable)                   │
│                                                                             │
│  acompletion() | aembedding() | get_model_info()                            │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│  OpenAI  │  Anthropic  │  AWS Bedrock  │  Gemini  │  Ollama  │  100+ more   │
└─────────────────────────────────────────────────────────────────────────────┘

Key Design Decisions:

1. Static class methods - No instantiation; call LLMClient.create_chat_completion() directly
2. Standardized responses - ChatCompletionResponse and EmbeddingResponse dataclasses
3. Protocol-based testing - LLMBackend Protocol for mock injection via set_backend()
4. Exception hierarchy - LLMClientError, ModelValidationError, APIKeyMissingError
5. Model config resolution - MODEL_CONFIGS → LiteLLM get_model_info() → fallback

---

Changes

New:

• agent_memory_server/llm/ package:
  • client.py - LLMClient with create_chat_completion(), create_embedding(), create_embeddings(), get_model_config(), optimize_query()
  • types.py - ChatCompletionResponse, EmbeddingResponse, LLMBackend Protocol
  • exceptions.py - LLMClientError, ModelValidationError, APIKeyMissingError
  • __init__.py - Re-exports for clean imports

Removed:

• agent_memory_server/llms.py

Updated:

• main.py - Startup validation via LLMClient.get_model_config()
• extraction.py - LLMClient.create_chat_completion()
• summarization.py - LLMClient.create_chat_completion()
• memory_strategies.py - LLMClient.create_chat_completion()
• long_term_memory.py - LLMClient.create_embedding()
• api.py - LLMClient.get_model_config()
• vectorstore_factory.py - Delegates to LLMClient.create_embeddings()
• All affected test files

---

Migration Guide

# Before
from agent_memory_server.llms import get_model_client
client = await get_model_client(model_name)
response = await client.create_chat_completion(model=model_name, prompt=messages)
text = response.choices[0].message.content

# After
from agent_memory_server.llm import LLMClient
response = await LLMClient.create_chat_completion(model=model_name, messages=messages)
text = response.content

# Embeddings - Before
from langchain_openai import OpenAIEmbeddings
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# Embeddings - After
from agent_memory_server.llm import LLMClient
embeddings = LLMClient.create_embeddings()  # Uses configured provider

# Custom endpoints / gateways
response = await LLMClient.create_chat_completion(
    model="gpt-4o",
    messages=messages,
    api_base="https://your-gateway.example.com",
    api_key="your-gateway-key"
)

---

Breaking Changes

None. This is an internal refactor. External APIs remain unchanged.

---

Testing

• All tests pass
• Added tests/test_llm_client.py with unit tests for the new client

---

Future Work

Option: Move to RedisVL Directly

Currently, langchain_redis.RedisVectorStore owns embedding calls—it expects an Embeddings object and calls embed_query()/embed_documents() internally. Moving to RedisVL directly would give us full control:

# Current: LangChain controls embedding
vectorstore = RedisVectorStore(embeddings=LLMClient.create_embeddings(), ...)
vectorstore.aadd_documents(docs)  # LangChain embeds internally

# Future: We control embedding
vectors = await LLMClient.create_embedding(model, texts)
query = VectorQuery(vector=vectors.embeddings[0], ...)
results = await index.query(query)

Trade-offs:

Aspect	LiteLLM (LLMClient)	RedisVL Vectorizers
Providers	100+	~6 (OpenAI, Azure, Cohere, HuggingFace, VertexAI, Mistral)
Already in codebase	Yes	New dependency path
Unified interface	Single LLMClient	Separate vectorizer classes
Async support	Native	Varies

Recommendation: If we migrate to RedisVL, continue using LLMClient.create_embedding() rather than RedisVL vectorizers—broader provider support and the abstraction is already built.

[nkanu17]

We use langchain_redis.RedisVectorStore for the embedding calls. It expects the langchain Embeddings object and calls embed_query()/embed_documents() internally.

If we move away from LangChain VectorStore to RedisVL directly, we can use either LiteLLM (via LLMClient.create_embedding()) or RedisVL's built-in vectorizers for embeddings. Since the repo now has LLMClient built as a unified abstraction with native async support, it would be nice to stick with LiteLLM. It gives broader provider coverage without introducing another embedding interface into the codebase.

For now this is not implemented.

[Copilot]

Pull request overview

This PR replaces custom LLM wrapper classes with a unified LLMClient abstraction layer backed by LiteLLM, significantly reducing maintenance burden and enabling support for multiple providers without code changes.

Key Changes:

• Introduced new agent_memory_server/llm/ package with LLMClient facade class
• Removed agent_memory_server/llms.py (454 lines) and replaced with cleaner abstraction
• Updated all LLM-consuming modules to use LLMClient.create_chat_completion() and LLMClient.create_embedding()

Reviewed changes

Copilot reviewed 67 out of 75 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
agent_memory_server/llm/client.py	New LLMClient facade with chat completion, embedding, and query optimization methods
agent_memory_server/llm/types.py	Standardized response dataclasses (ChatCompletionResponse, EmbeddingResponse)
agent_memory_server/llm/exceptions.py	Custom exception hierarchy for LLM operations
agent_memory_server/llm/__init__.py	Clean re-exports for package consumers
agent_memory_server/main.py	Updated model validation to use LLMClient.get_model_config()
agent_memory_server/summarization.py	Removed client parameter, uses LLMClient directly
agent_memory_server/memory_strategies.py	Updated to use LLMClient.create_chat_completion()
agent_memory_server/long_term_memory.py	Removed llm_client parameters throughout
agent_memory_server/vectorstore_factory.py	Delegates embedding creation to LLMClient.create_embeddings()
agent_memory_server/extraction.py	Updated to use LLMClient.create_chat_completion()
agent_memory_server/api.py	Uses LLMClient.get_model_config() for model configuration
tests/test_llm_client.py	New comprehensive test suite for LLMClient
uv.lock, pyproject.toml	Added litellm>=1.80.11 dependency
Multiple test files	Updated mocks to use ChatCompletionResponse dataclass
Java client files	Whitespace cleanup (trailing whitespace removal)
Documentation files	Updated references from llms.py to llm package

The PR is well-structured and follows good software engineering practices with proper abstraction, testing, and documentation updates. No critical issues found.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[vishal-bala]

I'd be curious to know your thinking behind opting for the Client/Backend abstraction as opposed to a general Client abstract base class (that has a LiteLLM implementation subclass). I think both approaches can work here, but I feel like the single global LLMClient object with a reconfigurable LLMBackend comes with some potential Python scoping gotchas.

This was more of a higher-level look through - I saw that the LLMClient currently is directly calling LiteLLM. Is that intentional and meant to be abstracted in the future, or actually part of the design?

[vishal-bala]

If LiteLLM is supposed to be the backend to the agnostic LLMClient wouldn't we rather have a generic call to an LLMBackend method here rather than locking in LiteLLM?

[nkanu17]

LLMClient (facade)
    ├── _backend: LLMBackend | None  # Only set in tests
    │
    ├── create_chat_completion()
    │       if _backend: use it (tests)
    │       else: call litellm.acompletion() directly
    │
    └── create_embedding()
            if _backend: use it (tests)
            else: call litellm.aembedding() directly

The _backend is only for tests, we wouldnt swap LiteLLM for something else. The protocol lets tests inject mocks without hitting real APIs.

I did consider what you said, we'd only want a different backend if we wanted to swap LiteLLM entirely for something else like:

• Direct SDK calls (openai, anthropic, boto3)
• A different unified library/custom gateway client

[nkanu17]

And my thought was to revisit this when we need it. I did not want to create placeholder abstraction without a need yet.

[nkanu17]

Update: I removed the _backend -> _test_handler/LLMBackend Protocol and replaced it with standard mocks.

[nkanu17]

I'd be curious to know your thinking behind opting for the Client/Backend abstraction as opposed to a general Client abstract base class (that has a LiteLLM implementation subclass). I think both approaches can work here, but I feel like the single global LLMClient object with a reconfigurable LLMBackend comes with some potential Python scoping gotchas.

This was more of a higher-level look through - I saw that the LLMClient currently is directly calling LiteLLM. Is that intentional and meant to be abstracted in the future, or actually part of the design?

@vishal-bala
Yea so the backend is required for tests only, allowing tests to mock LLM calls without hitting real APIs. You are also right that the _backend state could cause issues in concurrent tests. However, we mitigate this by:

• Only setting _backend in tests, never in production
• Using LLMClient.reset() in test teardown

Open to an alternative.

---

Actually, this was a good discussion - I realized the _backend and LLMBackend names are confusing. I am changing it to TestLLMHandler (which is the protocol tests should follow to mock) and _backend -> _test_handler

---

Update:

I removed the _backend -> _test_handler/LLMBackend Protocol and replaced it with standard mocks.

[abrookins]

Looks amazing! LGTM 👍