small fixes, mostly ggml updates

5 files changed, 79 insertions, 46 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 4e7a7cc7..5d8e7f4e 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -507,7 +507,8 @@ class Autopilot(ContinueBaseModel):
 
             async def create_title():
                 title = await self.continue_sdk.models.medium.complete(
-                    f'Give a short title to describe the current chat session. Do not put quotes around the title. The first message was: "{user_input}". The title is: '
+                    f'Give a short title to describe the current chat session. Do not put quotes around the title. The first message was: "{user_input}". Do not use more than 10 words. The title is: ',
+                    max_tokens=20,
                 )
                 title = remove_quotes_and_escapes(title)
                 self.session_info = SessionInfo(
diff --git a/continuedev/src/continuedev/core/context.py b/continuedev/src/continuedev/core/context.py
index d51a32e2..ffe22d63 100644
--- a/continuedev/src/continuedev/core/context.py
+++ b/continuedev/src/continuedev/core/context.py
@@ -257,6 +257,8 @@ class ContextManager:
                 await asyncio.wait_for(add_docs(), timeout=5)
             except asyncio.TimeoutError:
                 logger.warning("Failed to add document to meilisearch in 5 seconds")
+            except Exception as e:
+                logger.warning(f"Error adding document to meilisearch: {e}")
 
     @staticmethod
     async def delete_documents(ids):
diff --git a/continuedev/src/continuedev/libs/llm/ggml.py b/continuedev/src/continuedev/libs/llm/ggml.py
index f2c53e7b..34c3ab74 100644
--- a/continuedev/src/continuedev/libs/llm/ggml.py
+++ b/continuedev/src/continuedev/libs/llm/ggml.py
@@ -1,11 +1,16 @@
 import json
-from typing import Any, Coroutine, Dict, Generator, List, Union
+from typing import Any, Callable, Coroutine, Dict, Generator, List, Optional, Union
 
 import aiohttp
 
 from ...core.main import ChatMessage
 from ..llm import LLM
-from ..util.count_tokens import DEFAULT_ARGS, compile_chat_messages, count_tokens
+from ..util.count_tokens import (
+    DEFAULT_ARGS,
+    compile_chat_messages,
+    count_tokens,
+    format_chat_messages,
+)
 
 
 class GGML(LLM):
@@ -14,15 +19,15 @@ class GGML(LLM):
     server_url: str = "http://localhost:8000"
     verify_ssl: bool = True
 
-    _client_session: aiohttp.ClientSession = None
+    requires_write_log = True
+
+    write_log: Optional[Callable[[str], None]] = None
 
     class Config:
         arbitrary_types_allowed = True
 
-    async def start(self, **kwargs):
-        self._client_session = aiohttp.ClientSession(
-            connector=aiohttp.TCPConnector(verify_ssl=self.verify_ssl)
-        )
+    async def start(self, write_log: Callable[[str], None], **kwargs):
+        self.write_log = write_log
 
     async def stop(self):
         await self._client_session.close()
@@ -60,15 +65,24 @@ class GGML(LLM):
             system_message=self.system_message,
         )
 
-        async with self._client_session.post(
-            f"{self.server_url}/v1/completions", json={"messages": messages, **args}
-        ) as resp:
-            async for line in resp.content.iter_any():
-                if line:
-                    try:
-                        yield line.decode("utf-8")
-                    except:
-                        raise Exception(str(line))
+        self.write_log(f"Prompt: \n\n{format_chat_messages(messages)}")
+        completion = ""
+        async with aiohttp.ClientSession(
+            connector=aiohttp.TCPConnector(verify_ssl=self.verify_ssl)
+        ) as client_session:
+            async with client_session.post(
+                f"{self.server_url}/v1/completions", json={"messages": messages, **args}
+            ) as resp:
+                async for line in resp.content.iter_any():
+                    if line:
+                        try:
+                            chunk = line.decode("utf-8")
+                            yield chunk
+                            completion += chunk
+                        except:
+                            raise Exception(str(line))
+
+        self.write_log(f"Completion: \n\n{completion}")
 
     async def stream_chat(
         self, messages: List[ChatMessage] = None, **kwargs
@@ -86,31 +100,42 @@ class GGML(LLM):
         args["stream"] = True
 
         async def generator():
-            async with self._client_session.post(
-                f"{self.server_url}/v1/chat/completions",
-                json={"messages": messages, **args},
-            ) as resp:
-                # This is streaming application/json instaed of text/event-stream
-                async for line, end in resp.content.iter_chunks():
-                    json_chunk = line.decode("utf-8")
-                    if json_chunk.startswith(": ping - ") or json_chunk.startswith(
-                        "data: [DONE]"
-                    ):
-                        continue
-                    chunks = json_chunk.split("\n")
-                    for chunk in chunks:
-                        if chunk.strip() != "":
-                            yield json.loads(chunk[6:])["choices"][0][
-                                "delta"
-                            ]  # {"role": "assistant", "content": "..."}
+            async with aiohttp.ClientSession(
+                connector=aiohttp.TCPConnector(verify_ssl=self.verify_ssl)
+            ) as client_session:
+                async with client_session.post(
+                    f"{self.server_url}/v1/chat/completions",
+                    json={"messages": messages, **args},
+                ) as resp:
+                    # This is streaming application/json instaed of text/event-stream
+                    async for line, end in resp.content.iter_chunks():
+                        json_chunk = line.decode("utf-8")
+                        if json_chunk.startswith(": ping - ") or json_chunk.startswith(
+                            "data: [DONE]"
+                        ):
+                            continue
+                        chunks = json_chunk.split("\n")
+                        for chunk in chunks:
+                            if chunk.strip() != "":
+                                yield json.loads(chunk[6:])["choices"][0][
+                                    "delta"
+                                ]  # {"role": "assistant", "content": "..."}
 
         # Because quite often the first attempt fails, and it works thereafter
+        self.write_log(f"Prompt: \n\n{format_chat_messages(messages)}")
+        completion = ""
         try:
             async for chunk in generator():
                 yield chunk
+                if "content" in chunk:
+                    completion += chunk["content"]
         except:
             async for chunk in generator():
                 yield chunk
+                if "content" in chunk:
+                    completion += chunk["content"]
+
+        self.write_log(f"Completion: \n\n{completion}")
 
     async def complete(
         self, prompt: str, with_history: List[ChatMessage] = None, **kwargs
@@ -127,12 +152,18 @@ class GGML(LLM):
         #     system_message=self.system_message,
         # )
 
-        async with self._client_session.post(
-            f"{self.server_url}/v1/completions",
-            json={
-                "prompt": prompt,
-                **args,
-            },
-        ) as resp:
-            text = await resp.text()
-            return json.loads(text)["choices"][0]["text"]
+        self.write_log(f"Prompt: \n\n{prompt}")
+        async with aiohttp.ClientSession(
+            connector=aiohttp.TCPConnector(verify_ssl=self.verify_ssl)
+        ) as client_session:
+            async with client_session.post(
+                f"{self.server_url}/v1/completions",
+                json={
+                    "prompt": prompt,
+                    **args,
+                },
+            ) as resp:
+                text = await resp.text()
+                completion = json.loads(text)["choices"][0]["text"]
+                self.write_log(f"Completion: \n\n{completion}")
+                return completion
diff --git a/continuedev/src/continuedev/libs/util/commonregex.py b/continuedev/src/continuedev/libs/util/commonregex.py
index 3c4fb38c..9f119122 100644
--- a/continuedev/src/continuedev/libs/util/commonregex.py
+++ b/continuedev/src/continuedev/libs/util/commonregex.py
@@ -57,7 +57,6 @@ regexes = {
     "times": time,
     "phones": phone,
     "phones_with_exts": phones_with_exts,
-    "links": link,
     "emails": email,
     "ips": ip,
     "ipv6s": ipv6,
@@ -78,7 +77,6 @@ placeholders = {
     "times": "<TIME>",
     "phones": "<PHONE>",
     "phones_with_exts": "<PHONE_WITH_EXT>",
-    "links": "<LINK>",
     "emails": "<EMAIL>",
     "ips": "<IP>",
     "ipv6s": "<IPV6>",
diff --git a/continuedev/src/continuedev/plugins/steps/chat.py b/continuedev/src/continuedev/plugins/steps/chat.py
index 63548698..ad09f193 100644
--- a/continuedev/src/continuedev/plugins/steps/chat.py
+++ b/continuedev/src/continuedev/plugins/steps/chat.py
@@ -95,7 +95,8 @@ class SimpleChatStep(Step):
 
         self.name = remove_quotes_and_escapes(
             await sdk.models.medium.complete(
-                f"{self.description}\n\nHere is a short title for the above chat message:"
+                f"{self.description}\n\nHere is a short title for the above chat message (no more than 10 words):",
+                max_tokens=20,
             )
         )