Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Commit d2d7710

Browse files
committed
Pass in copilot extra headers when analyzing output packages
When querying Copilot, the LLM requires certain headers to be explicitly set to accept a request. We were already parsing them and storing in context and even using in the input pipeline, but forgot to use them in the output pipeline to extract packages from the inbound (LLM-to-client) snippets. Fixes: #335
1 parent 7154d67 commit d2d7710

File tree

1 file changed

+8
-7
lines changed
  • src/codegate/pipeline/extract_snippets

1 file changed

+8
-7
lines changed

src/codegate/pipeline/extract_snippets/output.py

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from litellm.types.utils import Delta, StreamingChoices
66

77
from codegate.llm_utils.extractor import PackageExtractor
8-
from codegate.pipeline.base import CodeSnippet, PipelineContext, PipelineSensitiveData
8+
from codegate.pipeline.base import CodeSnippet, PipelineContext
99
from codegate.pipeline.extract_snippets.extract_snippets import extract_snippets
1010
from codegate.pipeline.output import OutputPipelineContext, OutputPipelineStep
1111
from codegate.storage import StorageEngine
@@ -39,14 +39,15 @@ def _create_chunk(self, original_chunk: ModelResponse, content: str) -> ModelRes
3939
object="chat.completion.chunk",
4040
)
4141

42-
async def _snippet_comment(self, snippet: CodeSnippet, secrets: PipelineSensitiveData) -> str:
42+
async def _snippet_comment(self, snippet: CodeSnippet, context: PipelineContext) -> str:
4343
"""Create a comment for a snippet"""
4444
snippet.libraries = await PackageExtractor.extract_packages(
4545
content=snippet.code,
46-
provider=secrets.provider,
47-
model=secrets.model,
48-
api_key=secrets.api_key,
49-
base_url=secrets.api_base,
46+
provider=context.sensitive.provider if context.sensitive else None,
47+
model=context.sensitive.model if context.sensitive else None,
48+
api_key=context.sensitive.api_key if context.sensitive else None,
49+
base_url=context.sensitive.api_base if context.sensitive else None,
50+
extra_headers=context.metadata.get("extra_headers", None),
5051
)
5152
# Check if any of the snippet libraries is a bad package
5253
storage_engine = StorageEngine()
@@ -126,7 +127,7 @@ async def process_chunk(
126127
chunks.append(self._create_chunk(chunk, before))
127128
complete_comment += before
128129

129-
comment = await self._snippet_comment(last_snippet, input_context.sensitive)
130+
comment = await self._snippet_comment(last_snippet, input_context)
130131
complete_comment += comment
131132
chunks.append(
132133
self._create_chunk(

0 commit comments

Comments
 (0)