From b8da0d592165f40c8079c4170ea749085a046d36 Mon Sep 17 00:00:00 2001 From: Chris Trevino Date: Fri, 30 Oct 2026 13:48:20 -0700 Subject: [PATCH 1/2] Add optional Copilot SDK completion provider Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../minor-20261030203828497653.json | 4 + docs/config/models.md | 44 ++ packages/graphrag-llm/README.md | 113 ++++ .../completion/completion_factory.py | 7 + .../completion/copilot_completion.py | 492 +++++++++++++++ .../graphrag-llm/graphrag_llm/config/types.py | 1 + packages/graphrag-llm/pyproject.toml | 3 + pyproject.toml | 1 + .../language_model/test_copilot_completion.py | 571 ++++++++++++++++++ uv.lock | 22 + 10 files changed, 1258 insertions(+) create mode 100644 .semversioner/next-release/minor-20261030203828497653.json create mode 100644 packages/graphrag-llm/graphrag_llm/completion/copilot_completion.py create mode 100644 tests/unit/language_model/test_copilot_completion.py diff --git a/.semversioner/next-release/minor-20261030203828497653.json b/.semversioner/next-release/minor-20261030203828497653.json new file mode 100644 index 0000000000..980143d7f3 --- /dev/null +++ b/.semversioner/next-release/minor-20261030203828497653.json @@ -0,0 +1,4 @@ +{ + "type": "minor", + "description": "Add an optional GitHub Copilot SDK completion provider to graphrag-llm." +} diff --git a/docs/config/models.md b/docs/config/models.md index 401cba0318..2d39fdef6b 100644 --- a/docs/config/models.md +++ b/docs/config/models.md @@ -28,6 +28,50 @@ embedding_models: See [Detailed Configuration](yaml.md) for more details on configuration. [View LiteLLM basic usage](https://docs.litellm.ai/docs/#basic-usage) for details on how models are called (The `model_provider` is the portion prior to `/` while the `model` is the portion following the `/`). +## GitHub Copilot + +The built-in `copilot` completion provider uses the official GitHub Copilot SDK +and is registered through the same completion factory as other built-in +providers. It works with both GraphRAG's CLI and library APIs; no application-side +registration or provider conditionals are needed. + +Install the extra into the environment running GraphRAG (Python 3.11+): + +```shell +pip install "graphrag-llm[copilot]" +``` + +Authenticate with the Copilot CLI or the SDK's environment-based GitHub token +authentication, and choose an exact model ID available to your account: + +```yaml +completion_models: + default_completion_model: + type: copilot + model_provider: copilot + model: gpt-6.1-sol + timeout: 120 +``` + +Omit `api_key`, API endpoint settings, and unsupported `call_args` such as +`temperature`, `top_p`, `seed`, `max_tokens`, and `max_completion_tokens`. +The provider rejects unsupported parameters instead of silently ignoring them. +Copilot does not provide embeddings: keep `embedding_models` configured for a +different provider. Available models, rate limits and usage charges depend on +your Copilot subscription. + +DRIFT search is currently incompatible because it supplies sampling and +token-limit parameters internally, even when they are omitted from `call_args`. +Use another completion provider for DRIFT. + +The provider supports text completions, sync/async streaming and native +schema-based structured responses with Pydantic validation. Tools are disabled. +Requests use fresh sessions with no shared history; supplied multi-turn history +is rendered as a role-labelled transcript, not native role-separated messages. +Each uncached request starts and stops its own CLI client, which adds startup +overhead. See the [standalone provider documentation](https://github.com/microsoft/graphrag/tree/main/packages/graphrag-llm#github-copilot-completion-provider) +for lifecycle, authentication, usage metrics and other limitations. + ## Model Selection Considerations GraphRAG has been most thoroughly tested with the gpt-4 series of models from OpenAI, including gpt-4 gpt-4-turbo, gpt-4o, and gpt-4o-mini. Our [arXiv paper](https://arxiv.org/abs/2404.16130), for example, performed quality evaluation using gpt-4-turbo. As stated above, non-OpenAI models are supported through the use of LiteLLM but the suite of gpt-4 series of models from OpenAI remain the most tested and supported suite of models for GraphRAG – in other words, these are the models we know best and can help resolve issues with. diff --git a/packages/graphrag-llm/README.md b/packages/graphrag-llm/README.md index 5e3a1f1e53..6039525e9e 100644 --- a/packages/graphrag-llm/README.md +++ b/packages/graphrag-llm/README.md @@ -16,3 +16,116 @@ This examples demonstrates how to generate text embeddings using the GraphRAG LL [Open the notebook to explore the basic embeddings example code](example_notebooks/basic_embedding_example.ipynb) View the [notebooks](notebooks/README.md) for more examples. + +## GitHub Copilot completion provider + +The built-in `copilot` completion type uses the official +[GitHub Copilot SDK](https://github.com/github/copilot-sdk). It is available +independently of GraphRAG and requires Python 3.11 or later: + +```shell +pip install "graphrag-llm[copilot]" +``` + +Authenticate with the Copilot CLI before using the provider, or supply a GitHub +token through the SDK's environment-based authentication. You need Copilot access +to the requested model. No OpenAI/Azure API endpoint or API key is required. +Other providers do not require the Copilot extra. + +```python +from graphrag_llm.completion import create_completion +from graphrag_llm.config import ModelConfig +from graphrag_llm.utils import gather_completion_response + +completion = create_completion( + ModelConfig( + type="copilot", + model_provider="copilot", + model="gpt-6.1-sol", # Use an exact model ID available to your Copilot account. + timeout=120, + ) +) + +response = completion.completion(messages="What is retrieval-augmented generation?") +print(gather_completion_response(response)) + +stream = completion.completion(messages="Explain graph-based retrieval.", stream=True) +print(gather_completion_response(stream)) +``` + +Use `await completion.completion_async(...)` in asynchronous code; this returns +an async iterator when `stream=True`. Do not call the synchronous API inside a +running event loop. Fully consume streams, or explicitly call `close()` on a +synchronous stream / `await aclose()` on an asynchronous stream when stopping +early, so the request's resources are released immediately. + +Structured responses use the SDK's native per-request schema and the library's +normal Pydantic validation: + +```python +from pydantic import BaseModel + + +class Answer(BaseModel): + explanation: str + + +response = completion.completion( + messages="Explain graph-based retrieval in one sentence.", + response_format=Answer, +) +print(response.formatted_response) +``` + +Malformed output raises a parsing/validation error rather than being repaired or +reported as a successful structured completion. As with the LiteLLM provider, +`response_format` cannot be combined with streaming. The +`response_format_json_object=True` call option uses a JSON-only instruction and +validates the returned JSON object for non-streaming requests. Unlike typed +`response_format`, it is not native strict-schema enforcement. + +Provider-specific `ModelConfig` fields are `cli_path` (optional compatible CLI +executable), `github_token` (optional explicit token; prefer environment-based +authentication to storing credentials in configuration), and `timeout` (positive +seconds, default 600). `call_args` supports `timeout` and `reasoning_effort` on +models that support it; per-call values override these defaults. + +### Capabilities and limitations + +- Both APIs support text, streaming, caching, retries, rate limiting, and metrics + through the existing middleware. SDK-reported token usage is returned when + available; it is not an estimate of Copilot billing or premium requests. + Streaming usage is available on the final chunk with + `stream_options={"include_usage": True}`. As with other providers, streams + bypass the cache, middleware metrics describe stream creation rather than + completion, and errors during consumption are not retried automatically. +- Each uncached request starts a client and a fresh session. The provider checks + model availability, disables tools, repository instruction/config discovery, + file hooks, skills, memory and cross-session storage, and replaces the system prompt. + It disconnects and deletes its own session, then stops the client on success, + failure, cancellation, or explicit stream closure. This avoids persistent + background runtimes but adds startup overhead, especially for large batches. +- System/developer messages form the session system prompt. A single user + message is sent verbatim; multiple user/assistant turns are serialized as a + role-labelled JSON transcript in the prompt. This preserves supplied history + text but is not native OpenAI role-separated history or resumable SDK sessions. + No conversation state is shared between calls. +- Image/audio inputs, tool calling, embeddings, multiple choices (`n != 1`), + sampling parameters (`temperature`, `top_p`, `seed`, etc.) and token-limit + parameters are unsupported and raise errors rather than being silently + ignored. Configure a separate embedding provider. Local tokenization uses the + standard library tokenizer and is only an estimate for Copilot models. + GraphRAG's DRIFT search currently injects unsupported sampling parameters + internally and therefore requires another completion provider. +- Copilot model availability, quotas and schema support depend on your + subscription, selected model and compatible CLI runtime. The optional SDK + extra does not extend the Copilot provider to Python 3.10. Typed structured + responses use OpenAI-compatible strict schemas (including required properties + and closed objects); schemas unsupported by the selected model raise errors. + +**SDK runtime compatibility:** validation with SDK 1.0.26 and its default runtime +1.0.103 encountered an upstream event-decoding `AssertionError` for +`session.indexed_search` with `disabledReason=server_limit`, although the +completion succeeded. A compatible installed CLI selected with `cli_path` +completed structured and streaming checks without this callback error. The +provider does not suppress SDK callback errors or patch SDK internals. diff --git a/packages/graphrag-llm/graphrag_llm/completion/completion_factory.py b/packages/graphrag-llm/graphrag_llm/completion/completion_factory.py index e1923b9c63..20bcbdddda 100644 --- a/packages/graphrag-llm/graphrag_llm/completion/completion_factory.py +++ b/packages/graphrag-llm/graphrag_llm/completion/completion_factory.py @@ -104,6 +104,13 @@ def create_completion( completion_type=LLMProviderType.MockLLM, completion_initializer=MockLLMCompletion, ) + case LLMProviderType.Copilot: + from graphrag_llm.completion.copilot_completion import CopilotCompletion + + register_completion( + completion_type=LLMProviderType.Copilot, + completion_initializer=CopilotCompletion, + ) case _: msg = f"ModelConfig.type '{strategy}' is not registered in the CompletionFactory. Registered strategies: {', '.join(completion_factory.keys())}" raise ValueError(msg) diff --git a/packages/graphrag-llm/graphrag_llm/completion/copilot_completion.py b/packages/graphrag-llm/graphrag_llm/completion/copilot_completion.py new file mode 100644 index 0000000000..135e257e0e --- /dev/null +++ b/packages/graphrag-llm/graphrag_llm/completion/copilot_completion.py @@ -0,0 +1,492 @@ +# Copyright (c) 2024 Microsoft Corporation. +# Licensed under the MIT License + +"""Text completions through the optional GitHub Copilot SDK.""" + +import asyncio +import json +import math +import time +from collections.abc import AsyncGenerator, AsyncIterator, Iterator +from contextlib import aclosing, contextmanager, suppress +from typing import TYPE_CHECKING, Any, Unpack +from uuid import uuid4 + +from openai import pydantic_function_tool + +from graphrag_llm.completion.completion import LLMCompletion +from graphrag_llm.config.types import AuthMethod +from graphrag_llm.middleware import with_middleware_pipeline +from graphrag_llm.types import ( + LLMChoice, + LLMChoiceChunk, + LLMChoiceDelta, + LLMCompletionChunk, + LLMCompletionMessage, + LLMCompletionResponse, + LLMCompletionUsage, +) +from graphrag_llm.utils import structure_completion_response + +try: + from copilot import CopilotClient, RuntimeConnection + from copilot.generated.rpc import PermissionDecisionDeniedInteractivelyByUser + from copilot.session_events import ( + AssistantMessageData, + AssistantMessageDeltaData, + AssistantUsageData, + SessionEvent, + ) +except ModuleNotFoundError as exc: + if exc.name != "copilot": + raise + msg = ( + "The Copilot provider requires Python 3.11+ and the optional SDK. " + "Install it with: pip install 'graphrag-llm[copilot]'" + ) + raise ImportError(msg) from exc + +if TYPE_CHECKING: + from graphrag_cache import Cache, CacheKeyCreator + + from graphrag_llm.config import ModelConfig + from graphrag_llm.metrics import MetricsProcessor, MetricsStore + from graphrag_llm.rate_limit import RateLimiter + from graphrag_llm.retry import Retry + from graphrag_llm.tokenizer import Tokenizer + from graphrag_llm.types import ( + LLMCompletionArgs, + LLMCompletionMessagesParam, + Metrics, + ResponseFormat, + ) + + +def _deny_permission(*_args: Any) -> PermissionDecisionDeniedInteractivelyByUser: + return PermissionDecisionDeniedInteractivelyByUser( + feedback="Tools are disabled for GraphRAG LLM completions." + ) + + +def _text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list) and all( + isinstance(part, dict) + and part.get("type") == "text" + and isinstance(part.get("text"), str) + for part in content + ): + return "\n".join(part["text"] for part in content) + msg = "The Copilot provider supports text message content only." + raise ValueError(msg) + + +def _messages(messages: "LLMCompletionMessagesParam") -> tuple[str, str]: + if isinstance(messages, str): + return "You are a helpful assistant.", messages + system: list[str] = [] + turns: list[dict[str, str]] = [] + for message in messages: + role = message.get("role") + if ( + role not in {"system", "developer", "user", "assistant"} + or message.get("tool_calls") + or message.get("function_call") + ): + msg = "The Copilot provider supports system, developer, user and assistant text messages, without tool calls." + raise ValueError(msg) + content = _text(message.get("content")) + if role in {"system", "developer"}: + system.append(content) + else: + turns.append({"role": role, "content": content}) + if not turns: + msg = "At least one user or assistant message is required." + raise ValueError(msg) + prompt = turns[0]["content"] + if len(turns) != 1 or turns[0]["role"] != "user": + # The SDK accepts a prompt, not an OpenAI messages array. + prompt = ( + "Continue this conversation. Return only the next assistant message.\n" + f"Conversation (JSON):\n{json.dumps(turns)}" + ) + return "\n\n".join(system) or "You are a helpful assistant.", prompt + + +def _check_sync_context() -> None: + try: + asyncio.get_running_loop() + except RuntimeError: + pass + else: + msg = "Use completion_async inside a running event loop." + raise RuntimeError(msg) + + +@contextmanager +def _sync_loop() -> Iterator[asyncio.AbstractEventLoop]: + _check_sync_context() + loop = asyncio.new_event_loop() + try: + yield loop + finally: + try: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + finally: + loop.close() + + +class CopilotCompletion(LLMCompletion): + """Copilot text completion with request-scoped clients and tool-free sessions.""" + + def __init__( + self, + *, + model_id: str, + model_config: "ModelConfig", + tokenizer: "Tokenizer", + metrics_store: "MetricsStore", + cache_key_creator: "CacheKeyCreator", + metrics_processor: "MetricsProcessor | None" = None, + rate_limiter: "RateLimiter | None" = None, + retrier: "Retry | None" = None, + cache: "Cache | None" = None, + cli_path: str | None = None, + github_token: str | None = None, + timeout: float = 600, + **kwargs: Any, + ) -> None: + """Initialize the provider; CLI processes start only on a cache miss.""" + if kwargs: + msg = f"Unsupported Copilot configuration: {', '.join(sorted(kwargs))}" + raise ValueError(msg) + if ( + model_config.api_key + or model_config.api_base + or model_config.api_version + or model_config.azure_deployment_name + or model_config.auth_method != AuthMethod.ApiKey + ): + msg = "Copilot uses CLI authentication or github_token, not API endpoint settings." + raise ValueError(msg) + self._model = model_config.model + self._tokenizer = tokenizer + self._metrics_store = metrics_store + self._track_metrics = metrics_processor is not None + self._call_args = model_config.call_args + self._cli_path = cli_path + self._github_token = github_token + self._timeout = timeout + self._completion, self._completion_async = with_middleware_pipeline( + model_config=model_config, + model_fn=self._base_completion, + async_model_fn=self._base_completion_async, + request_type="chat", + cache=cache, + cache_key_creator=cache_key_creator, + tokenizer=tokenizer, + metrics_processor=metrics_processor, + rate_limiter=rate_limiter, + retrier=retrier, + ) + + def _prepare( + self, kwargs: dict[str, Any] + ) -> tuple[dict[str, Any], "Metrics | None"]: + args = {**self._call_args, **kwargs} + metrics = args.pop("metrics", None) + metrics = ( + (metrics if metrics is not None else {}) if self._track_metrics else None + ) + supported = { + "messages", + "stream", + "response_format", + "response_format_json_object", + "timeout", + "reasoning_effort", + "stream_options", + "n", + } + unsupported = [ + key + for key, value in args.items() + if key not in supported and value is not None + ] + if unsupported or args.get("n") not in (None, 1): + msg = f"Unsupported Copilot completion parameters: {', '.join(sorted(unsupported or ['n']))}" + raise ValueError(msg) + args = {key: value for key, value in args.items() if key in supported} + args.pop("n", None) + if args.get("stream") and args.get("response_format") is not None: + msg = "response_format is not supported for streaming completions." + raise ValueError(msg) + options = args.get("stream_options") or {} + if options.keys() - {"include_usage"}: + msg = "Copilot stream_options only supports include_usage." + raise ValueError(msg) + timeout = args.get("timeout") + args["timeout"] = self._timeout if timeout is None else timeout + if not math.isfinite(args["timeout"]) or args["timeout"] <= 0: + msg = "Copilot timeout must be a positive, finite number." + raise ValueError(msg) + # Validate before cache lookup, retries, or starting a runtime. + _messages(args["messages"]) + args["model"] = f"copilot/{self._model}" + return args, metrics + + def completion( + self, /, **kwargs: Unpack["LLMCompletionArgs[ResponseFormat]"] + ) -> "LLMCompletionResponse[ResponseFormat] | Iterator[LLMCompletionChunk]": + """Complete synchronously, including streaming and structured responses.""" + _check_sync_context() + args, metrics = self._prepare(dict(kwargs)) + try: + response = self._completion(metrics=metrics, **args) + if isinstance(response, LLMCompletionResponse) and args.get( + "response_format" + ): + response.formatted_response = structure_completion_response( + response.content, args["response_format"] + ) + return response + finally: + if metrics is not None: + self._metrics_store.update_metrics(metrics=metrics) + + async def completion_async( + self, /, **kwargs: Unpack["LLMCompletionArgs[ResponseFormat]"] + ) -> "LLMCompletionResponse[ResponseFormat] | AsyncIterator[LLMCompletionChunk]": + """Complete asynchronously, including streaming and structured responses.""" + args, metrics = self._prepare(dict(kwargs)) + try: + response = await self._completion_async(metrics=metrics, **args) + if isinstance(response, LLMCompletionResponse) and args.get( + "response_format" + ): + response.formatted_response = structure_completion_response( + response.content, args["response_format"] + ) + return response + finally: + if metrics is not None: + self._metrics_store.update_metrics(metrics=metrics) + + def _base_completion( + self, **kwargs: Any + ) -> LLMCompletionResponse | Iterator[LLMCompletionChunk]: + if kwargs.get("stream"): + return self._sync_stream(kwargs) + with _sync_loop() as loop: + return loop.run_until_complete(self._response(kwargs)) + + def _sync_stream(self, args: dict[str, Any]) -> Iterator[LLMCompletionChunk]: + with _sync_loop() as loop: + stream = self._stream(args) + try: + while True: + try: + chunk = loop.run_until_complete(anext(stream)) + except StopAsyncIteration: + return + yield chunk + finally: + loop.run_until_complete(stream.aclose()) + + async def _base_completion_async( + self, **kwargs: Any + ) -> LLMCompletionResponse | AsyncIterator[LLMCompletionChunk]: + if kwargs.get("stream"): + return self._stream(kwargs) + return await self._response(kwargs) + + async def _events(self, args: dict[str, Any]) -> AsyncGenerator[SessionEvent, None]: + system, prompt = _messages(args["messages"]) + schema = None + if args.get("response_format") is not None: + # The SDK forwards schemas without OpenAI's required strict normalization. + schema = pydantic_function_tool(args["response_format"])["function"].get( + "parameters" + ) + if schema is None: + msg = "Unable to derive a strict JSON schema from response_format." + raise ValueError(msg) + elif args.get("response_format_json_object"): + system += "\n\nReturn a single valid JSON object, without prose or Markdown fences." + client = CopilotClient( + connection=RuntimeConnection.for_stdio(path=self._cli_path), + github_token=self._github_token, + ) + try: + await asyncio.wait_for(client.start(), timeout=args["timeout"]) + models = await asyncio.wait_for( + client.list_models(), timeout=args["timeout"] + ) + if self._model not in {model.id for model in models}: + msg = f"Model '{self._model}' is not available through GitHub Copilot." + raise ValueError(msg) + session = await asyncio.wait_for( + client.create_session( + model=self._model, + system_message={"mode": "replace", "content": system}, + available_tools=[], + on_permission_request=_deny_permission, + skip_custom_instructions=True, + enable_config_discovery=False, + enable_on_demand_instruction_discovery=False, + skip_embedding_retrieval=True, + enable_file_hooks=False, + enable_skills=False, + enable_session_store=False, + mcp_servers={}, + custom_agents=[], + memory={"enabled": False}, + infinite_sessions={"enabled": False}, + streaming=bool(args.get("stream")), + reasoning_effort=args.get("reasoning_effort"), + ), + timeout=args["timeout"], + ) + try: + queue: asyncio.Queue[SessionEvent | None] = asyncio.Queue() + unsubscribe = session.on(queue.put_nowait) + task = asyncio.create_task( + session.send_and_wait( + prompt, response_schema=schema, timeout=args["timeout"] + ) + ) + task.add_done_callback(lambda _: queue.put_nowait(None)) + completed = False + try: + while (event := await queue.get()) is not None: + if not event.agent_id and isinstance( + event.data, (AssistantMessageDeltaData, AssistantUsageData) + ): + yield event + result = await task + if ( + result is None + or not isinstance(result.data, AssistantMessageData) + or result.data.tool_requests + or not result.data.content + ): + msg = "Copilot returned no final assistant text response." + raise RuntimeError(msg) + completed = True + yield result + finally: + unsubscribe() + if not task.done(): + task.cancel() + try: + with suppress(asyncio.CancelledError): + await task + finally: + if not completed: + await session.abort() + finally: + try: + await session.disconnect() + finally: + await client.delete_session(session.session_id) + finally: + await client.stop() + + async def _response(self, args: dict[str, Any]) -> LLMCompletionResponse: + usage: list[AssistantUsageData] = [] + content = "" + model = self._model + async with aclosing(self._events(args)) as events: + async for event in events: + if isinstance(event.data, AssistantUsageData): + usage.append(event.data) + elif isinstance(event.data, AssistantMessageData): + content = event.data.content + model = event.data.model or model + response = LLMCompletionResponse( + id=f"copilot-{uuid4()}", + object="chat.completion", + created=int(time.time()), + model=model, + choices=[ + LLMChoice( + index=0, + finish_reason="stop", + message=LLMCompletionMessage(role="assistant", content=content), + ) + ], + usage=_usage(usage) if usage else None, + ) + if args.get("response_format") is not None: + # Validate before caching; restore the concrete model after cache lookup. + structure_completion_response(content, args["response_format"]) + elif args.get("response_format_json_object") and not isinstance( + json.loads(content), dict + ): + msg = "Copilot did not return a JSON object." + raise ValueError(msg) + return response + + async def _stream( + self, args: dict[str, Any] + ) -> AsyncGenerator[LLMCompletionChunk, None]: + id = f"copilot-{uuid4()}" + created = int(time.time()) + usage: list[AssistantUsageData] = [] + emitted = False + + def chunk(content: str | None = None, *, final: bool = False): + return LLMCompletionChunk( + id=id, + object="chat.completion.chunk", + created=created, + model=self._model, + choices=[ + LLMChoiceChunk( + index=0, + delta=LLMChoiceDelta(content=content, role="assistant"), + finish_reason="stop" if final else None, + ) + ], + usage=_usage(usage) + if final + and usage + and (args.get("stream_options") or {}).get("include_usage") + else None, + ) + + events = self._events(args) + try: + async for event in events: + if isinstance(event.data, AssistantMessageDeltaData): + emitted = True + yield chunk(event.data.delta_content) + elif isinstance(event.data, AssistantUsageData): + usage.append(event.data) + elif isinstance(event.data, AssistantMessageData) and not emitted: + yield chunk(event.data.content) + finally: + await events.aclose() + yield chunk(final=True) + + @property + def metrics_store(self) -> "MetricsStore": + """The metrics store.""" + return self._metrics_store + + @property + def tokenizer(self) -> "Tokenizer": + """The tokenizer used for local token estimates.""" + return self._tokenizer + + +def _usage(records: list[AssistantUsageData]) -> LLMCompletionUsage: + prompt = sum(record.input_tokens or 0 for record in records) + completion = sum(record.output_tokens or 0 for record in records) + return LLMCompletionUsage( + prompt_tokens=prompt, + completion_tokens=completion, + total_tokens=prompt + completion, + ) diff --git a/packages/graphrag-llm/graphrag_llm/config/types.py b/packages/graphrag-llm/graphrag_llm/config/types.py index 320e8765fb..ad433b08d7 100644 --- a/packages/graphrag-llm/graphrag_llm/config/types.py +++ b/packages/graphrag-llm/graphrag_llm/config/types.py @@ -12,6 +12,7 @@ class LLMProviderType(StrEnum): LiteLLM = "litellm" MockLLM = "mock" + Copilot = "copilot" class AuthMethod(StrEnum): diff --git a/packages/graphrag-llm/pyproject.toml b/packages/graphrag-llm/pyproject.toml index 23b7cea8d8..6e32a6bee4 100644 --- a/packages/graphrag-llm/pyproject.toml +++ b/packages/graphrag-llm/pyproject.toml @@ -46,6 +46,9 @@ dependencies = [ "typing-extensions~=4.16", ] +[project.optional-dependencies] +copilot = ["github-copilot-sdk~=1.0.26; python_version >= '3.11'"] + [project.urls] Source = "https://github.com/microsoft/graphrag" diff --git a/pyproject.toml b/pyproject.toml index 9c32a6d888..8098d78b5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ requires-python = ">=3.11,<3.14" dev = [ "coverage~=7.16", "deptry~=0.25", + "github-copilot-sdk~=1.0.26", "ipykernel~=7.3", "jupyter~=1.1", "mkdocs-material~=9.7", diff --git a/tests/unit/language_model/test_copilot_completion.py b/tests/unit/language_model/test_copilot_completion.py new file mode 100644 index 0000000000..017e7ea599 --- /dev/null +++ b/tests/unit/language_model/test_copilot_completion.py @@ -0,0 +1,571 @@ +# Copyright (c) 2024 Microsoft Corporation. +# Licensed under the MIT License + +"""Copilot provider contract tests using the real SDK's types and signatures.""" + +import asyncio +import builtins +import importlib +import json +from collections.abc import AsyncGenerator, AsyncIterator, Generator, Iterator +from datetime import datetime, timezone +from typing import Any +from unittest.mock import create_autospec +from uuid import uuid4 + +import pytest +from copilot import CopilotClient +from copilot.session import CopilotSession +from copilot.session_events import ( + AssistantMessageData, + AssistantMessageDeltaData, + AssistantUsageData, + SessionEvent, + SessionEventType, +) +from graphrag_cache.memory_cache import MemoryCache +from graphrag_llm.completion import copilot_completion as provider +from graphrag_llm.completion import create_completion, register_completion +from graphrag_llm.config import ModelConfig +from graphrag_llm.config.retry_config import RetryConfig +from graphrag_llm.types import LLMCompletionResponse +from graphrag_llm.utils import ( + gather_completion_response, + gather_completion_response_async, +) +from pydantic import BaseModel, ValidationError + + +def event(data: Any, event_type: str, *, agent_id: str | None = None) -> SessionEvent: + return SessionEvent( + data=data, + type=SessionEventType(event_type), + id=uuid4(), + timestamp=datetime.now(timezone.utc), + agent_id=agent_id, + ) + + +@pytest.fixture +def sdk(monkeypatch: pytest.MonkeyPatch): + constructor = create_autospec(CopilotClient) + client = constructor.return_value + session = create_autospec(CopilotSession, instance=True) + session.session_id = "isolated-session" + client.create_session.return_value = session + client.list_models.return_value = [type("Model", (), {"id": "test-model"})()] + state: dict[str, Any] = { + "content": "Hello world", + "deltas": ["Hello", " world"], + "error": None, + "wait": False, + "no_response": False, + } + + def subscribe(callback): + state["callback"] = callback + return state.setdefault("unsubscribe", create_autospec(lambda: None)) + + session.on.side_effect = subscribe + + async def send(prompt, *, response_schema=None, timeout=60): + callback = state["callback"] + callback( + event( + AssistantUsageData( + model="test-model", + input_tokens=999, + output_tokens=999, + ), + "assistant.usage", + agent_id="sub-agent", + ) + ) + for delta in state["deltas"]: + callback( + event( + AssistantMessageDeltaData( + delta_content=delta, + message_id="response", + ), + "assistant.message_delta", + ) + ) + await asyncio.sleep(0) + if state["wait"]: + await asyncio.Event().wait() + if state["error"]: + raise state["error"] + callback( + event( + AssistantUsageData( + model="test-model", + input_tokens=10, + output_tokens=3, + ), + "assistant.usage", + ) + ) + if state["no_response"]: + return None + return event( + AssistantMessageData( + content=state["content"], + message_id="response", + model="test-model", + ), + "assistant.message", + ) + + session.send_and_wait.side_effect = send + monkeypatch.setattr(provider, "CopilotClient", constructor) + return constructor, client, session, state + + +def model(**kwargs: Any): + return create_completion( + ModelConfig( + type="copilot", + model_provider="copilot", + model="test-model", + **kwargs, + ) + ) + + +def assert_cleaned(sdk): + _, client, session, state = sdk + session.disconnect.assert_awaited_once() + client.delete_session.assert_awaited_once_with("isolated-session") + client.stop.assert_awaited_once() + state["unsubscribe"].assert_called_once() + + +def test_factory_sync_and_usage(sdk): + completion = model() + assert isinstance(completion, provider.CopilotCompletion) + response = completion.completion(messages="Hi") + assert isinstance(response, LLMCompletionResponse) + assert response.content == "Hello world" + assert response.model == "test-model" + assert response.usage + assert response.usage.total_tokens == 13 + assert completion.metrics_store.get_metrics()["total_tokens"] == 13 + assert completion.tokenizer.encode("hello") + _, client, session, _ = sdk + config = client.create_session.call_args.kwargs + assert config["available_tools"] == [] + assert config["skip_custom_instructions"] + assert not config["enable_config_discovery"] + assert not config["enable_on_demand_instruction_discovery"] + assert config["skip_embedding_retrieval"] + assert config["memory"] == {"enabled": False} + assert not config["enable_file_hooks"] + assert not config["enable_skills"] + assert not config["enable_session_store"] + assert config["infinite_sessions"] == {"enabled": False} + assert config["on_permission_request"](None, None).kind != "approved" + session.send_and_wait.assert_awaited_once_with( + "Hi", + response_schema=None, + timeout=600, + ) + assert_cleaned(sdk) + + +async def test_async_completion_and_conversation(sdk): + messages = [ + {"role": "system", "content": "Be concise"}, + {"role": "developer", "content": [{"type": "text", "text": "No tools"}]}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"}, + {"role": "user", "content": "Continue"}, + ] + completion = model(call_args={"reasoning_effort": "low", "timeout": 12}) + response = await completion.completion_async(messages=messages, timeout=4) + assert isinstance(response, LLMCompletionResponse) + _, client, session, _ = sdk + assert client.create_session.call_args.kwargs["system_message"] == { + "mode": "replace", + "content": "Be concise\n\nNo tools", + } + assert client.create_session.call_args.kwargs["reasoning_effort"] == "low" + prompt = session.send_and_wait.call_args.args[0] + assert json.loads(prompt.split("Conversation (JSON):\n")[1]) == messages[2:] + assert session.send_and_wait.call_args.kwargs["timeout"] == 4 + assert_cleaned(sdk) + + +class Answer(BaseModel): + answer: int + + +async def test_nested_schema_uses_strict_objects_and_required_defaults(sdk): + class NestedAnswer(BaseModel): + answer: Answer + explanation: str | None = None + + sdk[3]["content"] = '{"answer":{"answer":42},"explanation":null}' + response = await model().completion_async( + messages="Answer", response_format=NestedAnswer + ) + assert isinstance(response, LLMCompletionResponse) + assert response.formatted_response == NestedAnswer(answer=Answer(answer=42)) + schema = sdk[2].send_and_wait.call_args.kwargs["response_schema"] + assert schema["required"] == ["answer", "explanation"] + assert schema["additionalProperties"] is False + assert schema["$defs"]["Answer"]["additionalProperties"] is False + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_structured_response(sdk, asynchronous): + sdk[3]["content"] = '{"answer":42}' + completion = model() + if asynchronous: + response = await completion.completion_async( + messages="Answer", response_format=Answer + ) + else: + response = await asyncio.to_thread( + completion.completion, + messages="Answer", + response_format=Answer, + ) + assert isinstance(response, LLMCompletionResponse) + assert response.formatted_response == Answer(answer=42) + schema = sdk[2].send_and_wait.call_args.kwargs["response_schema"] + assert schema["properties"]["answer"]["type"] == "integer" + assert schema["required"] == ["answer"] + assert schema["additionalProperties"] is False + assert_cleaned(sdk) + + +@pytest.mark.parametrize("content", ['{"answer":"not a number"}', "not JSON"]) +async def test_invalid_structured_output_is_not_success(sdk, content): + sdk[3]["content"] = content + with pytest.raises((ValidationError, json.JSONDecodeError)): + await model().completion_async(messages="Answer", response_format=Answer) + assert_cleaned(sdk) + + +async def test_json_object(sdk): + sdk[3]["content"] = '{"answer":42}' + response = await model().completion_async( + messages="Answer", response_format_json_object=True + ) + assert isinstance(response, LLMCompletionResponse) + assert json.loads(response.content) == {"answer": 42} + assert sdk[2].send_and_wait.call_args.kwargs["response_schema"] is None + assert ( + "JSON object" + in sdk[1].create_session.call_args.kwargs["system_message"]["content"] + ) + + +@pytest.mark.parametrize("content", ["[]", "not JSON"]) +async def test_invalid_json_object_is_rejected(sdk, content): + sdk[3]["content"] = content + with pytest.raises((json.JSONDecodeError, ValueError), match=r"JSON|Expecting"): + await model().completion_async( + messages="Answer", response_format_json_object=True + ) + assert_cleaned(sdk) + + +def test_client_options_are_forwarded(sdk): + model(cli_path="custom-copilot", github_token="test-token").completion( + messages="Hi" + ) + assert sdk[0].call_args.kwargs["github_token"] == "test-token" + assert sdk[0].call_args.kwargs["connection"].path == "custom-copilot" + + +@pytest.mark.parametrize( + "kwargs", + [ + {"unsupported_option": True}, + {"api_key": "test-key"}, + {"api_base": "https://example.com"}, + {"azure_deployment_name": "deployment"}, + {"auth_method": "azure_managed_identity"}, + ], +) +def test_invalid_configuration_rejected(sdk, kwargs): + with pytest.raises(ValueError, match=r"Unsupported|CLI authentication"): + model(**kwargs) + sdk[0].assert_not_called() + + +def test_sync_call_preserves_caller_event_loop(sdk): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + model().completion(messages="Hi") + assert asyncio.get_event_loop() is loop + assert not loop.is_closed() + finally: + asyncio.set_event_loop(None) + loop.close() + + +def test_sync_streaming(sdk): + stream = model().completion( + messages="Hi", stream=True, stream_options={"include_usage": True} + ) + assert isinstance(stream, Iterator) + chunks = list(stream) + assert gather_completion_response(iter(chunks)) == "Hello world" + assert chunks[-1].choices[0].finish_reason == "stop" + assert chunks[-1].usage + assert chunks[-1].usage.total_tokens == 13 + assert len({chunk.id for chunk in chunks}) == 1 + assert_cleaned(sdk) + + +async def test_async_streaming(sdk): + stream = await model().completion_async(messages="Hi", stream=True) + assert isinstance(stream, AsyncIterator) + assert await gather_completion_response_async(stream) == "Hello world" + assert_cleaned(sdk) + + +async def test_async_stream_closed_early(sdk): + sdk[3]["wait"] = True + stream = await model().completion_async(messages="Hi", stream=True) + assert isinstance(stream, AsyncGenerator) + assert (await anext(stream)).choices[0].delta.content == "Hello" + await stream.aclose() + sdk[2].abort.assert_awaited_once() + assert_cleaned(sdk) + + +def test_sync_stream_closed_early(sdk): + sdk[3]["wait"] = True + stream = model().completion(messages="Hi", stream=True) + assert isinstance(stream, Generator) + assert next(stream).choices[0].delta.content == "Hello" + stream.close() + sdk[2].abort.assert_awaited_once() + assert_cleaned(sdk) + + +async def test_cancellation_cleans_up(sdk): + sdk[3]["wait"] = True + request = asyncio.create_task(model().completion_async(messages="Hi")) + while not sdk[2].send_and_wait.called: + await asyncio.sleep(0) + request.cancel() + with pytest.raises(asyncio.CancelledError): + await request + sdk[2].abort.assert_awaited_once() + assert_cleaned(sdk) + + +@pytest.mark.parametrize("streaming", [False, True]) +async def test_errors_propagate(sdk, streaming): + sdk[3]["error"] = TimeoutError("request timed out") + + async def consume(): + response = await model().completion_async(messages="Hi", stream=streaming) + if isinstance(response, AsyncIterator): + async for _ in response: + pass + + with pytest.raises(TimeoutError, match="request timed out"): + await consume() + sdk[2].abort.assert_awaited_once() + assert_cleaned(sdk) + + +async def test_no_response_is_error(sdk): + sdk[3]["no_response"] = True + with pytest.raises(RuntimeError, match="no final assistant"): + await model().completion_async(messages="Hi") + assert_cleaned(sdk) + + +async def test_unknown_model_does_not_fall_back(sdk): + sdk[1].list_models.return_value = [] + with pytest.raises(ValueError, match="not available"): + await model().completion_async(messages="Hi") + sdk[1].create_session.assert_not_called() + sdk[1].stop.assert_awaited_once() + + +@pytest.mark.parametrize("phase", ["start", "create_session"]) +async def test_startup_error_stops_client(sdk, phase): + getattr(sdk[1], phase).side_effect = RuntimeError("failed to start") + with pytest.raises(RuntimeError, match="failed to start"): + await model().completion_async(messages="Hi") + sdk[1].stop.assert_awaited_once() + + +async def test_disconnect_failure_still_deletes_and_stops(sdk): + sdk[2].disconnect.side_effect = RuntimeError("disconnect failed") + with pytest.raises(RuntimeError, match="disconnect failed"): + await model().completion_async(messages="Hi") + sdk[1].delete_session.assert_awaited_once() + sdk[1].stop.assert_awaited_once() + + +async def test_cache_hit_does_not_start_runtime(sdk): + completion = create_completion( + ModelConfig(type="copilot", model_provider="copilot", model="test-model"), + cache=MemoryCache(), + ) + first = await completion.completion_async(messages="Hi") + second = await completion.completion_async(messages="Hi") + assert isinstance(first, LLMCompletionResponse) + assert isinstance(second, LLMCompletionResponse) + assert first.content == second.content + sdk[0].assert_called_once() + + +async def test_cached_structured_response_keeps_pydantic_type(sdk): + sdk[3]["content"] = '{"answer":42}' + completion = create_completion( + ModelConfig(type="copilot", model_provider="copilot", model="test-model"), + cache=MemoryCache(), + ) + await completion.completion_async(messages="Answer", response_format=Answer) + cached = await completion.completion_async( + messages="Answer", response_format=Answer + ) + assert isinstance(cached, LLMCompletionResponse) + assert cached.formatted_response == Answer(answer=42) + sdk[0].assert_called_once() + + +async def test_cache_isolated_by_model(sdk): + cache = MemoryCache() + completion = create_completion( + ModelConfig(type="copilot", model_provider="copilot", model="test-model"), + cache=cache, + ) + await completion.completion_async(messages="Hi") + other = create_completion( + ModelConfig(type="copilot", model_provider="copilot", model="different-model"), + cache=cache, + ) + with pytest.raises(ValueError, match="not available"): + await other.completion_async(messages="Hi") + assert sdk[0].call_count == 2 + + +async def test_invalid_structured_output_is_not_cached(sdk): + sdk[3]["content"] = "not JSON" + completion = create_completion( + ModelConfig(type="copilot", model_provider="copilot", model="test-model"), + cache=MemoryCache(), + ) + with pytest.raises(json.JSONDecodeError): + await completion.completion_async(messages="Answer", response_format=Answer) + sdk[3]["content"] = '{"answer":42}' + response = await completion.completion_async( + messages="Answer", response_format=Answer + ) + assert isinstance(response, LLMCompletionResponse) + assert response.formatted_response == Answer(answer=42) + assert sdk[0].call_count == 2 + + +async def test_retry_uses_fresh_client(sdk): + sdk[1].start.side_effect = [RuntimeError("transient"), None] + response = await model( + retry=RetryConfig(type="immediate", max_retries=2) + ).completion_async(messages="Hi") + assert isinstance(response, LLMCompletionResponse) + assert sdk[0].call_count == 2 + assert sdk[1].stop.await_count == 2 + + +@pytest.mark.parametrize( + "kwargs", + [ + {"temperature": 0.1}, + {"tools": []}, + {"n": 2}, + {"stream": True, "response_format": Answer}, + {"timeout": 0}, + {"timeout": float("inf")}, + {"stream_options": {"invalid": True}}, + ], +) +async def test_unsupported_parameters_rejected_before_startup(sdk, kwargs): + with pytest.raises( + ValueError, match=r"Unsupported|not supported|timeout|stream_options" + ): + await model().completion_async(messages="Hi", **kwargs) + sdk[0].assert_not_called() + + +@pytest.mark.parametrize( + "messages", + [ + [], + [{"role": "system", "content": "Only system"}], + [{"role": "tool", "content": "tool output", "tool_call_id": "id"}], + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image"}, + } + ], + } + ], + [{"role": "assistant", "content": "", "tool_calls": [{"id": "tool"}]}], + ], +) +async def test_unsupported_messages_rejected(sdk, messages): + with pytest.raises(ValueError, match=r"supports|message is required"): + await model().completion_async(messages=messages) + sdk[0].assert_not_called() + + +async def test_sync_inside_event_loop_is_explicit_error(sdk): + await asyncio.sleep(0) + with pytest.raises(RuntimeError, match="completion_async"): + model().completion(messages="Hi") + sdk[0].assert_not_called() + + +async def test_cached_sync_call_inside_loop_is_explicit_error(sdk): + completion = create_completion( + ModelConfig(type="copilot", model_provider="copilot", model="test-model"), + cache=MemoryCache(), + ) + await asyncio.sleep(0) + with pytest.raises(RuntimeError, match="completion_async"): + completion.completion(messages="Hi") + sdk[0].assert_not_called() + + +def test_custom_registration_remains_supported(monkeypatch): + from graphrag_llm.completion.completion_factory import completion_factory + + initializer = create_autospec(provider.CopilotCompletion) + monkeypatch.setattr(completion_factory, "_service_initializers", {}) + register_completion("copilot", initializer) + completion = create_completion( + ModelConfig(type="copilot", model_provider="copilot", model="test-model") + ) + assert completion is initializer.return_value + + +def test_missing_optional_dependency_is_actionable(monkeypatch): + original_import = builtins.__import__ + + def missing_copilot(name, *args, **kwargs): + if name == "copilot": + msg = "missing copilot" + raise ModuleNotFoundError(msg, name="copilot") + return original_import(name, *args, **kwargs) + + with monkeypatch.context() as context: + context.setattr(builtins, "__import__", missing_copilot) + with pytest.raises(ImportError, match=r"graphrag-llm\[copilot\]"): + importlib.reload(provider) + importlib.reload(provider) diff --git a/uv.lock b/uv.lock index 05820f74d1..47deeb7d7b 100644 --- a/uv.lock +++ b/uv.lock @@ -1086,6 +1086,19 @@ wheels = [ { url = "https://ms-feed-25.pkgs.visualstudio.com/6f084628-a36d-42cb-934d-057357e379dc/_packaging/49d7402f-07bb-4b18-a9ce-086e6e98a554/pypi/download/ghp-import/2.1/ghp_import-2.1.0-py3-none-any.whl", hash = "sha256:8337dd7b50877f163d4c0289bc1f1c7f127550241988d568c1db512c4324a619", size = 11034, upload_time = "2026-07-28T03:16:18Z" }, ] +[[package]] +name = "github-copilot-sdk" +version = "1.0.26" +source = { registry = "https://packagefeedproxy.microsoft.io/pypi/simple" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, + { name = "python-dateutil" }, +] +wheels = [ + { url = "https://ms-feed-25.pkgs.visualstudio.com/6f084628-a36d-42cb-934d-057357e379dc/_packaging/49d7402f-07bb-4b18-a9ce-086e6e98a554/pypi/download/github-copilot-sdk/1.0.26/github_copilot_sdk-1.0.26-py3-none-any.whl", hash = "sha256:0854a5a34e3d89a8392ddafff5de78df2071937354f97c4865536b3c5ba15abe", size = 924388, upload_time = "2026-10-29T16:37:17Z" }, +] + [[package]] name = "graphrag" version = "3.3.0" @@ -1233,10 +1246,16 @@ dependencies = [ { name = "typing-extensions" }, ] +[package.optional-dependencies] +copilot = [ + { name = "github-copilot-sdk" }, +] + [package.metadata] requires-dist = [ { name = "azure-identity", specifier = "~=1.25" }, { name = "fastapi", specifier = ">=0.136.3,<1.0" }, + { name = "github-copilot-sdk", marker = "python_full_version >= '3.11' and extra == 'copilot'", specifier = "~=1.0.26" }, { name = "graphrag-cache", editable = "packages/graphrag-cache" }, { name = "graphrag-common", editable = "packages/graphrag-common" }, { name = "jinja2", specifier = "~=3.1" }, @@ -1246,6 +1265,7 @@ requires-dist = [ { name = "pydantic", specifier = "~=2.13" }, { name = "typing-extensions", specifier = "~=4.16" }, ] +provides-extras = ["copilot"] [[package]] name = "graphrag-monorepo" @@ -1256,6 +1276,7 @@ source = { virtual = "." } dev = [ { name = "coverage" }, { name = "deptry" }, + { name = "github-copilot-sdk" }, { name = "ipykernel" }, { name = "jupyter" }, { name = "mkdocs-exclude-search" }, @@ -1282,6 +1303,7 @@ dev = [ dev = [ { name = "coverage", specifier = "~=7.16" }, { name = "deptry", specifier = "~=0.25" }, + { name = "github-copilot-sdk", specifier = "~=1.0.26" }, { name = "ipykernel", specifier = "~=7.3" }, { name = "jupyter", specifier = "~=1.1" }, { name = "mkdocs-exclude-search", specifier = "~=0.6" }, From c87ae2fa4be68b7929b8252161b79609dd57f40b Mon Sep 17 00:00:00 2001 From: Chris Trevino Date: Fri, 30 Oct 2026 14:12:27 -0700 Subject: [PATCH 2/2] Fix Copilot documentation spellcheck Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- dictionary.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/dictionary.txt b/dictionary.txt index 75126c4640..0b682fa26c 100644 --- a/dictionary.txt +++ b/dictionary.txt @@ -5,6 +5,7 @@ Trinh Fernández # Pythonisms +aclose PYTHONPATH pycache nopython