Code editor

composer/pipeline/core.py

@@ -134,7 +134,7 @@ def _component_cache_key(c: ContractComponentInstance) -> CacheKey[Properties, C
     return CacheKey(string_hash("|".join([c.app.model_dump_json(), str(c.ind), str(c._contract.ind)])))
 
 
-def _batch_cache_key[FormT: BaseModel](props: list[PropertyFormulation]) -> CacheKey[ComponentGroup, FormT]:
+def _batch_cache_key[FormT: BaseModel](props: list[PropertyFormulation]) -> CacheKey[ComponentGroup, FormT]: # pyright: ignore[reportInvalidTypeVarUse]
     return CacheKey(string_hash("|".join(p.model_dump_json() for p in props)))
 
 

composer/spec/agent_index.py

@@ -3,10 +3,10 @@
 import os
 import unicodedata
 from dataclasses import dataclass
-from typing import TypedDict, Unpack, cast, overload, override
+from typing import TypedDict, Unpack, cast, overload, override, Awaitable, AsyncIterable
 from abc import abstractmethod, ABC
 
-from langgraph.store.base import BaseStore
+from langgraph.store.base import BaseStore, SearchItem
 from graphcore.tools.schemas import WithAsyncDependencies
 from pydantic import Field
 
@@ -79,8 +79,52 @@ def agent_index_config_from_env(data_ns: tuple[str, ...]) -> AgentIndexConfig:
         "Expected one of: tiered, trusted, readonly."
     )
 
+class AgentIndexBase:
+    @dataclass
+    class _ListIter[T]:
+        wrapped: list[T]
+        ptr: int = 0
+
+        def peek(self) -> T | None:
+            if self.ptr >= len(self.wrapped):
+                return None
+            return self.wrapped[self.ptr]
+
+        def pop(self) -> T:
+            to_ret = self.wrapped[self.ptr]
+            self.ptr += 1
+            return to_ret
+
+    @classmethod
+    def _normalize(cls, text: str) -> str:
+        nfkc = unicodedata.normalize("NFKC", text).casefold()
+        stripped = "".join(c for c in nfkc if not unicodedata.category(c).startswith("P"))
+        return " ".join(stripped.split())
 
-class AgentIndex:
+    @classmethod
+    def _question_key(
+        cls, question: str
+    ) -> str:
+        return hashlib.sha256(cls._normalize(question).encode()).hexdigest()[18:]
+
+    @classmethod
+    async def parallel_search(
+        cls, *args: Awaitable[list[SearchItem]]
+    ) -> AsyncIterable[SearchItem]:
+        query_results = await asyncio.gather(*args)
+        result_pointers = [
+            cls._ListIter(l) for l in query_results
+        ]
+        while True:
+            query = ((i, peeked) for (i, it) in enumerate(result_pointers) if (peeked := it.peek()) is not None)
+            m = max(query, key=lambda r: cast(float, r[1].score), default=None)
+            if m is None:
+                return
+            popped = result_pointers[m[0]].pop()
+            yield popped
+
+
+class AgentIndex(AgentIndexBase):
     """Two-layer semantic cache.
 
     ``base_layer`` is always consulted on reads and is the default
@@ -142,16 +186,6 @@ def _read_pools(self) -> list[tuple[str, ...]]:
             return [self.base_layer]
         return [self.write_layer, self.base_layer]
 
-    def _normalize(self, text: str) -> str:
-        nfkc = unicodedata.normalize("NFKC", text).casefold()
-        stripped = "".join(c for c in nfkc if not unicodedata.category(c).startswith("P"))
-        return " ".join(stripped.split())
-
-    def _question_key(
-        self, question: str
-    ) -> str:
-        return hashlib.sha256(self._normalize(question).encode()).hexdigest()[18:]
-
     async def aput(
         self,
         **doc: Unpack[AgentResult]
@@ -183,20 +217,13 @@ async def aget(
                 return cast(AgentResult, r.value)
         return None
 
-    @dataclass
-    class _ListIter[T]:
-        wrapped: list[T]
-        ptr: int = 0
-
-        def peek(self) -> T | None:
-            if self.ptr >= len(self.wrapped):
-                return None
-            return self.wrapped[self.ptr]
-
-        def pop(self) -> T:
-            to_ret = self.wrapped[self.ptr]
-            self.ptr += 1
-            return to_ret
+    def _raw_search(
+        self, question: str
+    ) -> list[Awaitable[list[SearchItem]]]:
+        return [
+            self.store.asearch(ns, query=question, limit=5)
+            for ns in self._read_pools
+        ]
 
     async def asearch(
         self, question: str
@@ -209,21 +236,12 @@ async def asearch(
         # the same metric (cosine similarity), so merging by score is
         # meaningful. Dedup defends against the same key existing in both
         # pools (which a manual / offline promotion may produce).
-        pool_results = await asyncio.gather(*[
-            self.store.asearch(ns, query=question, limit=5)
-            for ns in self._read_pools
-        ])
-        result_pointers = [
-            AgentIndex._ListIter(l) for l in pool_results
-        ]
         context : list[IndexedAgentResult] = []
         seen = set()
-        while True:
-            query = ((i, peeked) for (i, it) in enumerate(result_pointers) if (peeked := it.peek()) is not None)
-            m = max(query, key=lambda r: cast(float, r[1].score), default=None)
-            if m is None:
-                return context
-            popped = result_pointers[m[0]].pop()
+        async for popped in self.parallel_search(*(
+            self.store.asearch(ns, query=question, limit=5)
+            for ns in self._read_pools
+        )):
             if popped.key in seen:
                 continue
             seen.add(popped.key)
@@ -234,6 +252,7 @@ async def asearch(
             })
             if len(context) == 5:
                 return context
+        return context
 
     @overload
     @staticmethod
@@ -330,10 +349,10 @@ async def run(self) -> str:
                 # Document-Ref can be surfaced.
                 return answer
             return f"{answer}\n\nDocument-Ref: {ref_key}"
-        
+
 class RetrieveDocumentTool(WithAsyncDependencies[str, AgentIndex]):
     """
-    Retrieve the document associated with the provided document ref
+    Retrieve the document associated with the provided document ref.
     """
     ref: str = Field(description="The document reference id")
 

composer/spec/code_explorer.py

@@ -9,20 +9,17 @@
 
 from pydantic import Field, BaseModel
 
-from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.tools import BaseTool
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.checkpoint.memory import InMemorySaver
 
-from graphcore.graph import Builder, FlowInput, MessagesState
+from graphcore.graph import FlowInput, MessagesState
 from graphcore.tools.schemas import WithAsyncImplementation, WithInjectedId
-from graphcore.tools.vfs import fs_tools
 
 from composer.spec.graph_builder import bind_standard, run_to_completion
-from composer.templates.loader import load_jinja_template
 from composer.spec.tool_env import BaseSourceTools, BasicAgentTools
 from composer.spec.util import uniq_thread_id
-from composer.spec.agent_index import AgentIndex, IndexedTool, WithAgentIndex
+from composer.spec.agent_index import AgentIndex, IndexedTool
 from composer.ui.tool_display import tool_display_of, CommonTools
 
 
@@ -85,10 +82,19 @@ class _ExploreCodeCommon(BaseModel):
     is roughly the slowest single answer instead of the sum.
     """
     question: str = Field(
-        description="A specific, focused question about the source code. "
-        "Good: 'What state variables does withdraw() modify and how?' "
-        "Bad: 'Tell me about the contract' "
-        "Bad: 'What is the definition of function X?' (read the source directly)"
+        description="""
+A specific, focused question about the source code. Do not ask questions about:
+* The current task you're working on
+* How to use other tools
+* Questions about CVL or the prover
+* Protocol related questions unrelated to the source code (e.g. expected deployment params, contract address seeds)
+* Questions about the Solidity language itself
+
+Good: 'What state variables does withdraw() modify and how?'
+Bad: 'Tell me about the contract'
+Bad: 'What is the definition of function X?' (read the source directly)
+Bad: 'Is it realistic to expect deposits > 2^128?'
+"""
     )
 
 

composer/spec/context.py

@@ -121,9 +121,16 @@ class Contract:
 
 type Abstraction = CVLGeneration | FoundryGeneration
 
+class EditorAgent:
+    """editor for a single property"""
+
+class EditorJudge:
+    """judge for the editor"""
+
+
 type Marker = (
     InvJudge | InvFormal | Properties | ComponentGroup
-    | CVLJudge | FoundryJudge | Abstraction | Contract
+    | CVLJudge | FoundryJudge | Abstraction | Contract | EditorAgent | EditorJudge
 )
 
 # ---------------------------------------------------------------------------