fix(cody_chat): updating cody_chat to work with SG5.6.0

examples/python-5.6-cody-client/cody_chat.py

@@ -1,27 +1,35 @@
-import requests
-import os
+"""
+This module implements a Cody Chat CLI that interacts with the Sourcegraph API
+to provide context-aware responses based on repository content.
+"""
+
 import argparse
+import json
+import os
 
+import requests
+import sseclient
 
-access_token = os.getenv('SRC_ACCESS_TOKEN')
+access_token = os.getenv("SRC_ACCESS_TOKEN")
 if not access_token:
     raise ValueError("Error: SRC_ACCESS_TOKEN environment variable is not set.")
-endpoint = os.getenv('SRC_ENDPOINT')
+endpoint = os.getenv("SRC_ENDPOINT")
 if not endpoint:
     raise ValueError("Error: SRC_ENDPOINT environment variable is not set.")
 # GraphQL endpoint URL
 graphql_url = endpoint + "/.api/graphql"
-chat_completions_url = endpoint + "/.api/completions/stream?api-version=1&client-name=jetbrains&client-version=6.0.0-SNAPSHOT'"
+chat_completions_url = (
+    endpoint
+    + "/.api/completions/stream?api-version=1&client-name=jetbrains&client-version=6.0.0-SNAPSHOT'"
+)
 # Headers
-headers = {
-    "Content-Type": "application/json",
-    "Authorization": f"token {access_token}"
-}
+headers = {"Content-Type": "application/json", "Authorization": f"token {access_token}"}
+
 
 def format_context(context):
     """
     Format the context for the given repositories and query.
-    
+
     :param context: The context to format
     :return: The formatted context (string)
     """
@@ -30,7 +38,9 @@ def format_context(context):
     context_parts.append("<context>")
     for result in context:
         context_parts.append("<item>")
-        context_parts.append(f"<file>{result['blob']['path']}:{result['startLine']}-{result['endLine']}</file>")
+        context_parts.append(
+            f"<file>{result['blob']['path']}:{result['startLine']}-{result['endLine']}</file>"
+        )
         context_parts.append(f"<chunk>{result['chunkContent']}</chunk>")
         context_parts.append("</item>")
     context_parts.append("</context>")
@@ -39,19 +49,20 @@ def format_context(context):
 
     return formatted_context
 
+
 def get_repo_context(repo_names, query, code_results_count=10, text_results_count=5):
     """
     Get repository context using a GraphQL query.
-    
+
     :param repos: List of repository IDs
     :param query: Search query string
     :param code_results_count: Number of code results to return (default: 10)
     :param text_results_count: Number of text results to return (default: 5)
     :return: JSON response containing the repository context
     """
-    CONTEXT_SEARCH_QUERY = """
-    query GetCodyContext($repos: [ID!]!, $query: String!, $codeResultsCount: Int!, $textResultsCount: Int!, $filePatterns: [String!]) {
-        getCodyContext(repos: $repos, query: $query, codeResultsCount: $codeResultsCount, textResultsCount: $textResultsCount, filePatterns: $filePatterns) {
+    context_search_query = """
+    query GetCodyContext($repos: [ID!]!, $query: String!, $codeResultsCount: Int!, $textResultsCount: Int!) {
+        getCodyContext(repos: $repos, query: $query, codeResultsCount: $codeResultsCount, textResultsCount: $textResultsCount) {
             ...on FileChunkContext {
                 blob {
                     path
@@ -80,26 +91,30 @@ def get_repo_context(repo_names, query, code_results_count=10, text_results_coun
         "query": query,
         "codeResultsCount": code_results_count,
         "textResultsCount": text_results_count,
-        "filePatterns": []
     }
 
-    response = requests.post(graphql_url, json={"query": CONTEXT_SEARCH_QUERY, "variables": variables}, headers=headers)
-
+    response = requests.post(
+        graphql_url,
+        json={"query": context_search_query, "variables": variables},
+        headers=headers,
+        timeout=30,
+    )
+
     if response.status_code == 200:
-        return format_context(response.json()['data']['getCodyContext'])
-    else:
-        print(f"Request failed with status code: {response.status_code}")
-        return None
+        return format_context(response.json()["data"]["getCodyContext"])
+
+    print(f"Request failed with status code: {response.status_code}")
+    return None
 
 
 def get_repo_ids(repo_names):
     """
     Convert repository names to their corresponding IDs using a GraphQL query.
-    
+
     :param repo_names: List of repository names
     :return: Dictionary mapping repository names to their IDs
     """
-    REPOSITORY_IDS_QUERY = """
+    repository_ids_query = """
     query Repositories($names: [String!]!, $first: Int!) {
         repositories(names: $names, first: $first) {
             nodes {
@@ -109,62 +124,74 @@ def get_repo_ids(repo_names):
         }
     }
     """
-
-    variables = {
-        "names": repo_names,
-        "first": len(repo_names)
-    }
-
+
+    variables = {"names": repo_names, "first": len(repo_names)}
+
     response = requests.post(
         graphql_url,
-        json={"query": REPOSITORY_IDS_QUERY, "variables": variables},
-        headers=headers
+        json={"query": repository_ids_query, "variables": variables},
+        headers=headers,
+        timeout=30,
     )
-    
+
     if response.status_code == 200:
         data = response.json()
-        return {node['name']: node['id'] for node in data['data']['repositories']['nodes']}
-    else:
-        print(f"Failed to fetch repository IDs. Status code: {response.status_code}")
-        return {}
+        return {
+            node["name"]: node["id"] for node in data["data"]["repositories"]["nodes"]
+        }
+
+    print(f"Failed to fetch repository IDs. Status code: {response.status_code}")
+    return {}
 
 
 def chat_completions(query):
     """
-    Send a chat completion request to the Sourcegraph API.
+    Send a chat completion request to the Sourcegraph API and process the SSE stream.
 
     :param query: The user's query (string)
-    :return: The completion response (string)
+    :return: The last completion response (string)
     """
     data = {
-        'maxTokensToSample': 4000,
-        'messages': [
-            {
-                'speaker': 'human',
-                'text': query
-            }
-        ],
-        'model': 'gpt-4o',
-        'temperature': 0.2,
-        'topK': -1,
-        'topP': -1,
-        'stream': False
+        "maxTokensToSample": 4000,
+        "messages": [{"speaker": "human", "text": query}],
+        "model": "gpt-4o",
+        "temperature": 0.2,
+        "topK": -1,
+        "topP": -1,
+        "stream": True,
     }
 
-    response = requests.post(chat_completions_url, headers=headers, json=data)
-
-    if response.status_code == 200:
-        return response.json().get('completion', '')
-    else:
-        print(f"Request failed with status code: {response.status_code}")
+    try:
+        response = requests.post(
+            chat_completions_url, headers=headers, json=data, stream=True, timeout=30
+        )
+        response.raise_for_status()
+
+        client = sseclient.SSEClient(response)
+        last_response = ""
+        for event in client.events():
+            if event.data != "[DONE]":
+                try:
+                    event_data = json.loads(event.data)
+                    if "completion" in event_data:
+                        last_response = event_data["completion"]
+                except json.JSONDecodeError:
+                    print(f"Error decoding event data: {event.data}")
+
+        return last_response
+
+    except requests.exceptions.RequestException as e:
+        print(f"Request failed: {str(e)}")
+        if hasattr(e, "response") and e.response is not None:
+            print(f"Status code: {e.response.status_code}")
+            print(f"Response content: {e.response.text}")
         return None
 
 
-
 def cody_chat(repo_names, query):
     """
     Get context for the given repositories and query, then print it out.
-    
+
     :param repo_names: List of repository names (strings)
     :param query: Natural language query (string)
     """
@@ -183,13 +210,19 @@ def cody_chat(repo_names, query):
 
     response = chat_completions(final_prompt)
     print(response)
-
 
 
 def main():
+    """
+    Main function to handle command-line arguments and execute the Cody Chat CLI.
+    """
     parser = argparse.ArgumentParser(description="Cody Chat CLI")
-    parser.add_argument('--context-repo', action='append', help='Repository name (can be used multiple times)')
-    parser.add_argument('--message', type=str, help='Query message')
+    parser.add_argument(
+        "--context-repo",
+        action="append",
+        help="Repository name (can be used multiple times)",
+    )
+    parser.add_argument("--message", type=str, help="Query message")
 
     args = parser.parse_args()
 
@@ -203,5 +236,6 @@ def main():
     query = args.message
     cody_chat(repo_names, query)
 
+
 if __name__ == "__main__":
     main()

examples/python-5.6-cody-client/pyproject.toml

@@ -1,9 +1,13 @@
 [project]
 name = "python-cody-client"
 version = "0.1.0"
-description = "Add your description here"
+description = "Olaf Geirsson"
 readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "requests>=2.32.3",
+    "sseclient-py",
+    "black",
+    "isort",
+    "pylint",
 ]

examples/python-5.6-cody-client/requirements.txt

@@ -0,0 +1,12 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml
+certifi==2024.8.30
+    # via requests
+charset-normalizer==3.3.2
+    # via requests
+idna==3.10
+    # via requests
+requests==2.32.3
+    # via python-cody-client (pyproject.toml)
+urllib3==2.2.3
+    # via requests