Skip to content

Commit 2bd3922

Browse files
authored
Litellm auth fix (#2083)
* Fix scope for Azure auth with LiteLLM * Change internal language on max_attempts to max_retries * Rework model config connectivity validation * Semver * Swtich smoke tests to LiteLLM * Take out temporary retry_strategy = none since it is not fnllm compatible * Bump smoke test timeout * Bump smoke timeout further * Tune smoke params * Update smoke test bounds * Remove covariates from min-csv smoke * Smoke: adjust communities, remove drift * Remove secrets where they aren't necessary * Clean out old env var references
1 parent 7f996cf commit 2bd3922

File tree

24 files changed

+2144
-2367
lines changed

24 files changed

+2144
-2367
lines changed

.github/workflows/gh-pages.yml

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,6 @@ jobs:
1515
GH_PAGES: 1
1616
DEBUG: 1
1717
GRAPHRAG_API_KEY: ${{ secrets.GRAPHRAG_API_KEY }}
18-
GRAPHRAG_LLM_MODEL: ${{ secrets.GRAPHRAG_LLM_MODEL }}
19-
GRAPHRAG_EMBEDDING_MODEL: ${{ secrets.GRAPHRAG_EMBEDDING_MODEL }}
2018

2119
steps:
2220
- uses: actions/checkout@v4

.github/workflows/python-notebook-tests.yml

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,6 @@ jobs:
3838
env:
3939
DEBUG: 1
4040
GRAPHRAG_API_KEY: ${{ secrets.OPENAI_NOTEBOOK_KEY }}
41-
GRAPHRAG_LLM_MODEL: ${{ secrets.GRAPHRAG_LLM_MODEL }}
42-
GRAPHRAG_EMBEDDING_MODEL: ${{ secrets.GRAPHRAG_EMBEDDING_MODEL }}
4341

4442
runs-on: ${{ matrix.os }}
4543
steps:

.github/workflows/python-smoke-tests.yml

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -37,20 +37,8 @@ jobs:
3737
fail-fast: false # Continue running all jobs even if one fails
3838
env:
3939
DEBUG: 1
40-
GRAPHRAG_LLM_TYPE: "azure_openai_chat"
41-
GRAPHRAG_EMBEDDING_TYPE: "azure_openai_embedding"
4240
GRAPHRAG_API_KEY: ${{ secrets.OPENAI_API_KEY }}
4341
GRAPHRAG_API_BASE: ${{ secrets.GRAPHRAG_API_BASE }}
44-
GRAPHRAG_API_VERSION: ${{ secrets.GRAPHRAG_API_VERSION }}
45-
GRAPHRAG_LLM_DEPLOYMENT_NAME: ${{ secrets.GRAPHRAG_LLM_DEPLOYMENT_NAME }}
46-
GRAPHRAG_EMBEDDING_DEPLOYMENT_NAME: ${{ secrets.GRAPHRAG_EMBEDDING_DEPLOYMENT_NAME }}
47-
GRAPHRAG_LLM_MODEL: ${{ secrets.GRAPHRAG_LLM_MODEL }}
48-
GRAPHRAG_EMBEDDING_MODEL: ${{ secrets.GRAPHRAG_EMBEDDING_MODEL }}
49-
# We have Windows + Linux runners in 3.10, so we need to divide the rate limits by 2
50-
GRAPHRAG_LLM_TPM: 200_000 # 400_000 / 2
51-
GRAPHRAG_LLM_RPM: 1_000 # 2_000 / 2
52-
GRAPHRAG_EMBEDDING_TPM: 225_000 # 450_000 / 2
53-
GRAPHRAG_EMBEDDING_RPM: 1_000 # 2_000 / 2
5442
# Azure AI Search config
5543
AZURE_AI_SEARCH_URL_ENDPOINT: ${{ secrets.AZURE_AI_SEARCH_URL_ENDPOINT }}
5644
AZURE_AI_SEARCH_API_KEY: ${{ secrets.AZURE_AI_SEARCH_API_KEY }}
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "patch",
3+
"description": "Fix Azure auth scope issue with LiteLLM."
4+
}

DEVELOPING.md

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -119,8 +119,3 @@ and then in your bashrc, add
119119
Make sure you have python3.10-dev installed or more generally `python<version>-dev`
120120

121121
`sudo apt-get install python3.10-dev`
122-
123-
### LLM call constantly exceeds TPM, RPM or time limits
124-
125-
`GRAPHRAG_LLM_THREAD_COUNT` and `GRAPHRAG_EMBEDDING_THREAD_COUNT` are both set to 50 by default. You can modify these values
126-
to reduce concurrency. Please refer to the [Configuration Documents](https://microsoft.github.io/graphrag/config/overview/)

docs/developing.md

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,3 @@ Make sure llvm-9 and llvm-9-dev are installed:
7777
and then in your bashrc, add
7878

7979
`export LLVM_CONFIG=/usr/bin/llvm-config-9`
80-
81-
### LLM call constantly exceeds TPM, RPM or time limits
82-
83-
`GRAPHRAG_LLM_THREAD_COUNT` and `GRAPHRAG_EMBEDDING_THREAD_COUNT` are both set to 50 by default. You can modify these values
84-
to reduce concurrency. Please refer to the [Configuration Documents](config/overview.md)

docs/prompt_tuning/auto_prompt_tuning.md

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -79,15 +79,7 @@ After that, it uses one of the following selection methods to pick a sample to w
7979

8080
## Modify Env Vars
8181

82-
After running auto tuning, you should modify the following environment variables (or config variables) to pick up the new prompts on your index run. Note: Please make sure to update the correct path to the generated prompts, in this example we are using the default "prompts" path.
83-
84-
- `GRAPHRAG_ENTITY_EXTRACTION_PROMPT_FILE` = "prompts/entity_extraction.txt"
85-
86-
- `GRAPHRAG_COMMUNITY_REPORT_PROMPT_FILE` = "prompts/community_report.txt"
87-
88-
- `GRAPHRAG_SUMMARIZE_DESCRIPTIONS_PROMPT_FILE` = "prompts/summarize_descriptions.txt"
89-
90-
or in your yaml config file:
82+
After running auto tuning, you should modify the following config variables to pick up the new prompts on your index run. Note: Please make sure to update the correct path to the generated prompts, in this example we are using the default "prompts" path.
9183

9284
```yaml
9385
entity_extraction:

graphrag/config/models/graph_rag_config.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ def _validate_retry_services(self) -> None:
107107

108108
_ = retry_factory.create(
109109
strategy=model.retry_strategy,
110-
max_attempts=model.max_retries,
110+
max_retries=model.max_retries,
111111
max_retry_wait=model.max_retry_wait,
112112
)
113113

graphrag/index/validate_config.py

Lines changed: 36 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -15,42 +15,39 @@
1515

1616

1717
def validate_config_names(parameters: GraphRagConfig) -> None:
18-
"""Validate config file for LLM deployment name typos."""
19-
# Validate Chat LLM configs
20-
# TODO: Replace default_chat_model with a way to select the model
21-
default_llm_settings = parameters.get_language_model_config("default_chat_model")
22-
23-
llm = ModelManager().register_chat(
24-
name="test-llm",
25-
model_type=default_llm_settings.type,
26-
config=default_llm_settings,
27-
callbacks=NoopWorkflowCallbacks(),
28-
cache=None,
29-
)
30-
31-
try:
32-
asyncio.run(llm.achat("This is an LLM connectivity test. Say Hello World"))
33-
logger.info("LLM Config Params Validated")
34-
except Exception as e: # noqa: BLE001
35-
logger.error(f"LLM configuration error detected. Exiting...\n{e}") # noqa
36-
sys.exit(1)
37-
38-
# Validate Embeddings LLM configs
39-
embedding_llm_settings = parameters.get_language_model_config(
40-
parameters.embed_text.model_id
41-
)
42-
43-
embed_llm = ModelManager().register_embedding(
44-
name="test-embed-llm",
45-
model_type=embedding_llm_settings.type,
46-
config=embedding_llm_settings,
47-
callbacks=NoopWorkflowCallbacks(),
48-
cache=None,
49-
)
50-
51-
try:
52-
asyncio.run(embed_llm.aembed_batch(["This is an LLM Embedding Test String"]))
53-
logger.info("Embedding LLM Config Params Validated")
54-
except Exception as e: # noqa: BLE001
55-
logger.error(f"Embedding LLM configuration error detected. Exiting...\n{e}") # noqa
56-
sys.exit(1)
18+
"""Validate config file for model deployment name typos, by running a quick test message for each."""
19+
for id, config in parameters.models.items():
20+
if config.type in ["chat", "azure_openai", "openai"]:
21+
llm = ModelManager().register_chat(
22+
name="test-llm",
23+
model_type=config.type,
24+
config=config,
25+
callbacks=NoopWorkflowCallbacks(),
26+
cache=None,
27+
)
28+
try:
29+
asyncio.run(
30+
llm.achat("This is an LLM connectivity test. Say Hello World")
31+
)
32+
logger.info("LLM Config Params Validated")
33+
except Exception as e: # noqa: BLE001
34+
logger.error(f"LLM configuration error detected.\n{e}") # noqa
35+
print(f"Failed to validate language model ({id}) params", e) # noqa: T201
36+
sys.exit(1)
37+
elif config.type in ["embedding", "azure_openai_embedding", "openai_embedding"]:
38+
embed_llm = ModelManager().register_embedding(
39+
name="test-embed-llm",
40+
model_type=config.type,
41+
config=config,
42+
callbacks=NoopWorkflowCallbacks(),
43+
cache=None,
44+
)
45+
try:
46+
asyncio.run(
47+
embed_llm.aembed_batch(["This is an LLM Embedding Test String"])
48+
)
49+
logger.info("Embedding LLM Config Params Validated")
50+
except Exception as e: # noqa: BLE001
51+
logger.error(f"Embedding configuration error detected.\n{e}") # noqa
52+
print(f"Failed to validate embedding model ({id}) params", e) # noqa: T201
53+
sys.exit(1)

graphrag/language_model/providers/litellm/chat_model.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,9 +86,10 @@ def _create_base_completions(
8686
msg = "Azure Managed Identity authentication is only supported for Azure models."
8787
raise ValueError(msg)
8888

89+
base_args["azure_scope"] = base_args.pop("audience")
8990
base_args["azure_ad_token_provider"] = get_bearer_token_provider(
9091
DefaultAzureCredential(),
91-
COGNITIVE_SERVICES_AUDIENCE,
92+
model_config.audience or COGNITIVE_SERVICES_AUDIENCE,
9293
)
9394

9495
def _base_completion(**kwargs: Any) -> ModelResponse | CustomStreamWrapper:

0 commit comments

Comments
 (0)