Add executecode plugin

Author: codelion

optillm/plugins/executecode_plugin.py

@@ -0,0 +1,97 @@
+import re
+from typing import Tuple, List
+import nbformat
+from nbconvert.preprocessors import ExecutePreprocessor
+import os
+import tempfile
+import json
+
+SLUG = "executecode"
+
+def extract_python_code(text: str) -> List[str]:
+    """Extract Python code blocks from text."""
+    pattern = r'```python\s*(.*?)\s*```'
+    return re.findall(pattern, text, re.DOTALL)
+
+def execute_code(code: str) -> str:
+    """Execute Python code in a Jupyter notebook environment."""
+    with tempfile.NamedTemporaryFile(suffix='.ipynb', delete=False) as tmp:
+        notebook = nbformat.v4.new_notebook()
+        notebook['cells'] = [nbformat.v4.new_code_cell(code)]
+        nbformat.write(notebook, tmp)
+
+    try:
+        with open(tmp.name) as f:
+            nb = nbformat.read(f, as_version=4)
+        ep = ExecutePreprocessor(timeout=30, kernel_name='python3')
+        ep.preprocess(nb, {'metadata': {'path': './'}})
+
+        # Extract the output
+        output = ""
+        for cell in nb.cells:
+            if cell.cell_type == 'code' and cell.outputs:
+                for output_item in cell.outputs:
+                    if output_item.output_type == 'stream':
+                        output += output_item.text
+                    elif output_item.output_type == 'execute_result':
+                        output += str(output_item.data.get('text/plain', ''))
+        
+        return output.strip()
+    finally:
+        os.unlink(tmp.name)
+
+def should_execute_request_code(query: str) -> bool:
+    """Decide whether to execute code from the request based on the query."""
+    keywords = ['run', 'execute', 'output', 'result']
+    return any(keyword in query.lower() for keyword in keywords)
+
+def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str, int]:
+    query, request_code = extract_python_code(initial_query)[0] if extract_python_code(initial_query) else (initial_query, "")
+    
+    if should_execute_request_code(query) and request_code:
+        # Execute code from the request
+        code_output = execute_code(request_code)
+        context = f"Query: {query}\nCode:\n```python\n{request_code}\n```\nOutput:\n{code_output}"
+        
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": context}
+        ]
+        
+        response = client.chat.completions.create(
+            model=model,
+            messages=messages,
+        )
+        
+        return response.choices[0].message.content.strip(), response.usage.completion_tokens
+    else:
+        # Get initial response from the model
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": initial_query}
+        ]
+        
+        response = client.chat.completions.create(
+            model=model,
+            messages=messages,
+        )
+        
+        initial_response = response.choices[0].message.content.strip()
+        response_code = extract_python_code(initial_response)
+        
+        if response_code:
+            # Execute code from the response
+            code_output = execute_code(response_code[0])
+            context = f"Initial response:\n{initial_response}\n\nCode output:\n{code_output}"
+            
+            messages.append({"role": "assistant", "content": initial_response})
+            messages.append({"role": "user", "content": f"Based on the code execution output, please provide a final response:\n{context}"})
+            
+            final_response = client.chat.completions.create(
+                model=model,
+                messages=messages,
+            )
+            
+            return final_response.choices[0].message.content.strip(), response.usage.completion_tokens + final_response.usage.completion_tokens
+        else:
+            return initial_response, response.usage.completion_tokens

optillm/plugins/memory_plugin.py

@@ -20,7 +20,7 @@ def add(self, item: str):
         self.items.append(item)
         self.vectors = None  # Reset vectors to force recalculation
 
-    def get_relevant(self, query: str, n: int = 5) -> List[str]:
+    def get_relevant(self, query: str, n: int = 10) -> List[str]:
         if not self.items:
             return []
 
@@ -49,26 +49,46 @@ def extract_query(text: str) -> Tuple[str, str]:
             query = "What is the main point of this text?"
     return query, context
 
-def extract_key_information(text: str, client, model: str) -> List[str]:
-    # print(f"Prompt : {text}")
-    prompt = f"""Extract key information from the following text. Provide a list of important facts or concepts, each on a new line:
+def classify_margin(margin):
+        return margin.startswith("YES#")
 
+def extract_key_information(system_message, text: str, query: str, client, model: str) -> List[str]:
+    # print(f"Prompt : {text}")
+    messages = [
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": f"""
+'''text
 {text}
-
-Key information:"""
+'''
+Copy over all context relevant to the query: {query}
+Provide the answer in the format: <YES/NO>#<Relevant context>.
+Here are rules:
+- If you don't know how to answer the query - start your answer with NO#
+- If the text is not related to the query - start your answer with NO#
+- If you can extract relevant information - start your answer with YES#
+- If the text does not mention the person by name - start your answer with NO#
+Example answers:
+- YES#Western philosophy originated in Ancient Greece in the 6th century BCE with the pre-Socratics.
+- NO#No relevant context.
+"""}
+    ]
 
     try: 
         response = client.chat.completions.create(
             model=model,
-            messages=[{"role": "user", "content": prompt}],
+            messages=messages,
             max_tokens=1000
         )
-        key_info = response.choices[0].message.content.strip().split('\n')
+        key_info = response.choices[0].message.content.strip()
     except Exception as e:
         print(f"Error parsing content: {str(e)}")
         return [],0
+    margins = []
+
+    if classify_margin(key_info):
+        margins.append(key_info.split("#", 1)[1])
 
-    return [info.strip('- ') for info in key_info if info.strip()], response.usage.completion_tokens
+    return margins, response.usage.completion_tokens
 
 def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str, int]:
     memory = Memory()
@@ -80,7 +100,7 @@ def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str
     for i in range(0, len(context), chunk_size):
         chunk = context[i:i+chunk_size]
         # print(f"chunk: {chunk}")
-        key_info, tokens = extract_key_information(chunk, client, model)
+        key_info, tokens = extract_key_information(system_prompt, chunk, query, client, model)
         #print(f"key info: {key_info}")
         completion_tokens += tokens
         for info in key_info:
@@ -90,16 +110,22 @@ def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str
     relevant_info = memory.get_relevant(query)
     # print(f"relevant_info : {relevant_info}")
     # Generate response using relevant information
-    prompt = f"""System: {system_prompt}
-
-Context: {' '.join(relevant_info)}
-
+    messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"""
+
+I asked my assistant to read and analyse the above content page by page to help you complete this task. These are margin notes left on each page:
+'''text
+{relevant_info}
+'''
+Read again the note(s), take a deep breath and answer the query.
 {query}
-"""
+"""}
+    ]
 
     response = client.chat.completions.create(
         model=model,
-        messages=[{"role": "user", "content": prompt}],
+        messages=messages,
     )
     # print(f"response : {response}")
     final_response = response.choices[0].message.content.strip()

optillm/plugins/readurls_plugin.py

@@ -22,7 +22,7 @@ def extract_urls(text: str) -> List[str]:
 
     return cleaned_urls
 
-def fetch_webpage_content(url: str, max_length: int = 40000) -> str:
+def fetch_webpage_content(url: str, max_length: int = 100000) -> str:
     try:
         headers = {
             'User-Agent': 'optillm/0.0.1 (hhttps://github.com/codelion/optillm)'