open-context-engine 0.1.3 → 0.1.5

package/src/retrieval/evidence.py

@@ -0,0 +1,312 @@
+"""Evidence-oriented retrieval with per-question scoring and complete snippets.
+
+The historical engines remain available for frozen evaluations. This engine
+keeps independent recall channels, scores actual subquestions, and selects
+implementation evidence before supporting tests when that is what was asked.
+"""
+from collections import defaultdict
+import math
+import re
+import time
+
+import numpy as np
+
+from engine import Engine, document, post
+from reranker import rerank_pairs
+
+VERSION = 'evidence-v2'
+
+
+def source_role(unit):
+    path = unit['path'].lower()
+    if re.search(r'(^|/)(tests?|__tests__|fixtures)(/|$)|(^|/)test_[^/]+|[._](test|spec)\.', path):
+        return 'test'
+    if unit.get('language') == 'text' or re.search(r'(^|/)(docs?|examples?)/|\.(md|rst|txt)$', path):
+        return 'documentation'
+    return 'implementation'
+
+
+def requested_role(query):
+    if re.search(r'^(?:find|show|locate|list)\s+(?:the\s+)?(?:(?:unit|regression)\s+)?tests?\b|^(?:查找|找到|列出|展示).{0,5}(?:测试|用例)', query, re.I):
+        return 'test'
+    if re.search(r'^(?:find|show|locate|list)\s+(?:the\s+)?(?:docs?|documentation|readme)\b|^(?:查找|找到|列出|展示).{0,5}文档', query, re.I):
+        return 'documentation'
+    return 'implementation'
+
+
+def requested_languages(query):
+    """Honor explicit language scope without guessing from repository names."""
+    languages = set()
+    for pattern, values in [
+        (r'\bpython\b', {'python'}),
+        (r'\bgolang\b', {'go'}),
+        (r'\btypescript\b', {'typescript'}),
+        (r'\bjavascript\b', {'javascript', 'typescript'}),
+        (r'\b(?:frontend|front-end|react)\b|前端', {'javascript', 'typescript'}),
+    ]:
+        if re.search(pattern, query, re.I):
+            languages.update(values)
+    if re.search(r'\bGo\b|\bgo\s+(?:code|service|backend|worker|implementation)\b', query):
+        languages.add('go')
+    return languages
+
+
+class EvidenceEngine(Engine):
+    version = VERSION
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.roles = [source_role(u) for u in self.units]
+        self.symbols = defaultdict(list)
+        self.neighbors = [set() for _ in self.units]
+        self.callees = [set() for _ in self.units]
+        for u in self.units:
+            if u.get('name'):
+                self.symbols[u['symbol']].append(u['id'])
+            for relation in u.get('relations', []):
+                if relation['kind'] not in {'calls', 'same_symbol', 'references_value'}:
+                    continue
+                target = relation['target']
+                self.neighbors[u['id']].add(target)
+                self.neighbors[target].add(u['id'])
+                if relation['kind'] in {'calls', 'references_value'}:
+                    self.callees[u['id']].add(target)
+
+    def substance(self, uid):
+        unit = self.units[uid]
+        if self.roles[uid] != 'implementation':
+            return 1.0
+        if unit['kind'] in {'property', 'parameter', 'type', 'interface', 'type-alias'}:
+            return .25
+        lines = [line.strip() for line in unit['text'].splitlines() if line.strip()]
+        executable = [line for line in lines if not line.startswith(
+            ('//', '/*', '*', '#', 'import ', 'from ', 'export type ', '}'))]
+        if not executable:
+            return .1
+        if unit['kind'] in {'function', 'method'}:
+            return 1.0 if len(executable) >= 2 else .5
+        return 1.0 if len(executable) >= 5 else .25
+
+    def bundle(self, uid, budget):
+        unit = self.units[uid]
+        family = self.symbols.get(unit['symbol'], [uid]) if unit.get('name') else [uid]
+        if unit['kind'] not in {'function', 'method'}:
+            return [uid]
+        if sum(self.costs[i] for i in family) <= 3200:
+            return sorted(family, key=lambda i: self.units[i]['start'])
+        return [uid]
+
+    def render_selection(self, selected):
+        # Coalesce only returned source; never invent lines between snippets.
+        paths = {}
+        for uid in selected:
+            unit = self.units[uid]
+            lines = paths.setdefault(unit['path'], {})
+            for number, line in enumerate(unit['text'].split('\n'), unit['start']):
+                lines[number] = line
+        blocks = []
+        for path, lines in paths.items():
+            previous = None
+            for number in sorted(lines):
+                if previous is None or number != previous + 1:
+                    blocks.append('Path: ' + path)
+                blocks.append(f'{number}\t{lines[number]}')
+                previous = number
+            blocks.append('')
+        return '\n'.join(blocks)
+
+    def search(self, plan, budget=8000):
+        started = time.monotonic()
+        declarations = bool(re.search(r'\b(interface|type alias|schema)\b|类型定义|接口类型', plan['intent'], re.I))
+        def substance(uid):
+            if declarations and self.units[uid]['kind'] in {'type', 'interface', 'type-alias'}:
+                return 1.0
+            return self.substance(uid)
+        intent = plan['intent']
+        languages = requested_languages(intent)
+        scoped = np.asarray([not languages or u.get('language') in languages for u in self.units])
+        facets = list(dict.fromkeys(f['question'] for f in plan['facets'] if f['question'] != intent))[:4]
+        queries = [intent] + [facet + '\nContext for this part of the request: ' + intent for facet in facets]
+        result = post(self.embed_url + '/embeddings', {'model': self.embedding_model,
+            'input': ['Instruct: Retrieve source code implementing the requested behavior.\nQuery: ' + q for q in queries]}, self.embedding_key)
+        rows = sorted(result['data'], key=lambda row: row['index'])
+        if [row['index'] for row in rows] != list(range(len(queries))):
+            raise ValueError('Invalid query embedding indices')
+        matrix = np.asarray([row['embedding'] for row in rows], dtype=np.float32)
+        if matrix.shape != (len(queries), self.vectors.shape[1]) or not np.isfinite(matrix).all():
+            raise ValueError('Invalid query embedding vectors')
+        dense = self.vectors @ matrix.T
+        candidates, recall = set(), []
+        for col, query in enumerate(queries):
+            channels = []
+            for scores in (dense[:, col], self.lexical(query)):
+                order = np.argsort(-scores, kind='stable')
+                # Candidate depth is independent of the response token budget.
+                # Keep global recall and add a separate explicit-language lane.
+                ids = [int(uid) for uid in order[:40] if scores[uid] > 0]
+                if languages:
+                    local = [int(uid) for uid in order if scoped[uid] and scores[uid] > 0][:80]
+                    ids = list(dict.fromkeys(ids + local))
+                candidates.update(ids)
+                channels.append(ids)
+            recall.append(channels)
+        scores, waves = {}, []
+
+        def score(ids):
+            ids = sorted(ids)
+            pairs = [(q, j) for q in range(len(queries)) for j, uid in enumerate(ids) if (q, uid) not in scores]
+            if not pairs:
+                return
+            before = time.monotonic()
+            result = rerank_pairs(self.reranker, queries, [document(self.units[i], 5000) for i in ids], pairs, post)
+            for row in result['results']:
+                q, j = pairs[row['index']]
+                scores[q, ids[j]] = row['relevance_score']
+            waves.append({'pairs': len(pairs), 'requests': result['meta']['request_count'],
+                          'elapsedMs': round((time.monotonic() - before) * 1000)})
+
+        score(candidates)
+        primary = requested_role(plan['intent'])
+        seeds = set()
+        for q in range(len(queries)):
+            order = sorted(candidates, key=lambda uid: (-scores[q, uid], uid))
+            seeds.update(order[:8])
+            seeds.update([i for i in order if self.roles[i] == primary and substance(i) >= .5][:6])
+        support = defaultdict(float)
+        for uid in sorted(seeds):
+            outgoing = self.callees[uid]
+            neighbors = outgoing | self.neighbors[uid]
+            # Shared utility callers are broad hubs, not reliable task evidence.
+            if len(neighbors) > 24:
+                neighbors = outgoing
+            for neighbor in neighbors:
+                if neighbor not in candidates:
+                    confidence = max(scores[q, uid] for q in range(len(queries)))
+                    support[neighbor] = max(support[neighbor], confidence / max(1, len(neighbors)))
+        expanded = sorted(support, key=lambda uid: (-support[uid], uid))[:80]
+        candidates.update(expanded)
+        score(expanded)
+
+        completions = {i for uid in candidates for i in self.bundle(uid, budget)} - candidates
+        score(completions)
+        candidates.update(completions)
+
+        ranks = {}
+        for q in range(len(queries)):
+            for role in ('implementation', 'test', 'documentation'):
+                order = sorted((i for i in candidates if self.roles[i] == role),
+                               key=lambda uid: (-scores[q, uid], uid))
+                ranks.update({(q, uid): rank for rank, uid in enumerate(order)})
+
+        def value(q, uid):
+            relevance = scores[q, uid]
+            if relevance < .1:
+                return 0.0
+            # Rank separates saturated probability-like scores. No inverse cost
+            # reward: a short test must not displace a better-ranked body.
+            return (.5 * relevance + .5 * math.exp(-ranks[q, uid] / 8)) * substance(uid) * max(0.0, scores[0, uid]) * (1.0 if scoped[uid] else .25)
+
+        # Tests often describe behavior more explicitly than the implementation.
+        # Propagate relevance only along resolved source references, bounded by
+        # the test's fan-out. This is evidence support, not invented call edges.
+        corroboration = defaultdict(float)
+        for uid in candidates:
+            if self.roles[uid] != 'test':
+                continue
+            targets = [i for i in self.callees[uid] if i in candidates and self.roles[i] == 'implementation']
+            if not targets or len(targets) > 8:
+                continue
+            for q in range(len(queries)):
+                if scores[q, uid] < .1:
+                    continue
+                strength = (.5 * scores[q, uid] + .5 * math.exp(-ranks[q, uid] / 8)) / math.sqrt(len(targets))
+                for target in targets:
+                    corroboration[q, target] = max(corroboration[q, target], strength)
+        selected, selected_set, trace = [], set(), []
+        spent = 0
+        path_counts, covered = defaultdict(int), np.zeros(len(queries))
+        symbol_counts = defaultdict(int)
+
+        def choose(ids, q=None, phase='fill', limit=None):
+            nonlocal spent
+            choices = []
+            for uid in ids:
+                if uid in selected_set:
+                    continue
+                bundle = [uid]
+                if self.roles[uid] == primary == 'implementation':
+                    # Relevant local helpers are part of the entry point explanation.
+                    # Keep their source evidence with the calling implementation.
+                    references = {r['target'] for r in self.units[uid].get('relations', [])
+                                  if r.get('resolution') == 'explicit-doc-reference'}
+                    # Resolved local callees explain the selected entry point.
+                    # Bound fan-out, cost and relevance so utility hubs cannot
+                    # pull arbitrary dependencies into the answer.
+                    local_calls = {i for i in self.callees[uid]
+                                   if self.units[i]['path'] == self.units[uid]['path']}
+                    if len({self.units[i]['symbol'] for i in local_calls}) <= 8:
+                        references.update(local_calls)
+                    extras = [i for i in references if i in candidates and i not in selected_set and i != uid
+                              and self.roles[i] == 'implementation' and substance(i) >= .5
+                              and self.costs[i] <= 1000 and scores[0, i] >= .5]
+                    extras.sort(key=lambda i: (-scores[0, i], i))
+                    extra_cost = 0
+                    for i in extras[:3]:
+                        if extra_cost + self.costs[i] <= 1600:
+                            bundle.append(i)
+                            extra_cost += self.costs[i]
+                cost = sum(self.costs[i] for i in bundle)
+                if spent + cost > (budget if limit is None else limit):
+                    bundle, cost = [uid], self.costs[uid]
+                if spent + cost > (budget if limit is None else limit):
+                    continue
+                values = [value(col, uid) + .3 * corroboration[col, uid] for col in range(len(queries))]
+                merit = values[q] if q is not None else max(v / (1 + covered[col]) for col, v in enumerate(values))
+                merit /= 1 + .1 * path_counts[self.units[uid]['path']] + .8 * symbol_counts[self.units[uid]['symbol']]
+                if merit > 0:
+                    choices.append((merit, -uid, uid, bundle, cost, values))
+            if not choices:
+                return False
+            _, _, uid, bundle, cost, values = max(choices)
+            selected.extend(bundle); selected_set.update(bundle); spent += cost
+            covered[:] += values
+            path_counts[self.units[uid]['path']] += 1
+            symbol_counts[self.units[uid]['symbol']] += 1
+            trace.append({'id': uid, 'bundle': bundle, 'phase': phase, 'role': self.roles[uid], 'tokens': cost,
+                          'scores': [scores[q, uid] for q in range(len(queries))]})
+            return True
+
+        primary_ids = [i for i in candidates if self.roles[i] == primary and substance(i) >= .5]
+        facets = list(range(1, len(queries))) or [0]
+        # Build a stable 8k core, then add evidence. Candidate scoring and family
+        # completion do not depend on the output budget.
+        for stage in range(8000, max(8000, budget) + 8000, 8000):
+            ceiling = min(stage, budget)
+            for _ in range(3):
+                for q in facets:
+                    choose(primary_ids, q, 'primary', ceiling * .82)
+            for _ in range(2):
+                companions = {i for uid in selected for i in self.callees[uid]
+                              if i in candidates and self.roles[i] == primary and substance(i) >= .5}
+                choose(companions, phase='dependency', limit=ceiling * .82)
+            continuations = {i for uid in selected for i in self.bundle(uid, budget)}
+            while choose(continuations, phase='continuation', limit=ceiling * .85):
+                pass
+            for q in facets:
+                choose(primary_ids, q, 'primary', ceiling * .85)
+            if primary == 'implementation' and re.search(r'\btests?\b|测试|回归', intent, re.I):
+                tests = [i for i in candidates if self.roles[i] == 'test']
+                related = [i for i in tests if self.callees[i] & selected_set]
+                for q in facets:
+                    choose(related or tests, q, 'support', ceiling)
+            while choose(candidates, limit=ceiling):
+                pass
+        context = self.render_selection(selected)
+        tokens = len(self.encoding.encode_ordinary(context))
+        if tokens > budget:
+            raise ValueError('Evidence packing exceeded token budget')
+        return context, {'version': VERSION, 'elapsedMs': round((time.monotonic() - started) * 1000),
+            'tokens': tokens, 'queryCache': False, 'candidateCount': len(candidates),
+            'rerankedCount': len(candidates), 'expandedCount': len(expanded), 'queries': queries,
+            'modelRequests': {'embedding': 1, 'rerank': sum(w['requests'] for w in waves)},
+            'waves': waves, 'selected': trace, 'recall': recall}

package/src/retrieval/languages/__init__.py

@@ -8,7 +8,7 @@ import sys
 
 from . import python, typescript, text, go
 from .files import path_exclusion, read_text
-from .schema import SCHEMA_VERSION, SourceFile, validate_units
+from .schema import SCHEMA_VERSION, SourceFile, SourceSyntaxError, validate_units
 
 ADAPTERS = {'python': python, 'typescript': typescript, 'javascript': typescript, 'go': go, 'text': text}
 EXTENSIONS = {'.py': 'python', '.ts': 'typescript', '.tsx': 'typescript',
@@ -25,9 +25,12 @@ def adapter_manifest(files, language_options=None):
     options = language_options or {}
     if not isinstance(options, dict) or set(options) - set(ADAPTERS):
         raise ValueError('Unknown language options')
-    paths = [Path(__file__), Path(__file__).with_name('schema.py'), Path(__file__).with_name('files.py')]
+    paths = [Path(__file__), Path(__file__).with_name('schema.py'), Path(__file__).with_name('files.py'),
+             Path(text.__file__)]
     for language in languages:
         paths.append(Path(ADAPTERS[language].__file__))
+        if language == 'python':
+            paths.append(Path(python.__file__).with_name('python_calls.py'))
         if language in {'typescript', 'javascript'}:
             paths.append(Path(typescript.__file__).with_suffix('.mjs'))
         if language == 'go':
@@ -42,6 +45,29 @@ def adapter_manifest(files, language_options=None):
             'sourceSha256': {path.name: hashlib.sha256(path.read_bytes()).hexdigest() for path in paths}}
 
 
+def extract_with_fallback(adapter, sources, max_lines, settings):
+    try:
+        return adapter.extract(sources, max_lines, settings)
+    except SourceSyntaxError as error:
+        diagnostics = {item['path']: item for item in error.diagnostics}
+        if not diagnostics or not set(diagnostics) <= {source.path for source in sources}:
+            raise ValueError('Invalid syntax diagnostic paths') from error
+        # Rebuild relations using only valid source. Never preserve partial AST
+        # output or let a broken module supply symbols to its healthy callers.
+        valid = [source for source in sources if source.path not in diagnostics]
+        units = adapter.extract(valid, max_lines, settings) if valid else []
+        for source in sources:
+            if source.path not in diagnostics:
+                continue
+            fallback = text.extract([source], max_lines)
+            for unit in fallback:
+                unit['id'] += len(units)
+            if fallback:
+                fallback[0]['parseDiagnostic'] = {**diagnostics[source.path], 'fallback': 'text'}
+            units.extend(fallback)
+        return units
+
+
 def source_units(root, files, max_lines=65, language_options=None, report=None, cache=None):
     if type(max_lines) is not int or max_lines < 1:
         raise ValueError('max_lines must be a positive integer')
@@ -98,7 +124,7 @@ def source_units(root, files, max_lines=65, language_options=None, report=None,
             ], sort_keys=True).encode()).hexdigest() if cache is not None else None)
             previous = cache.get(key) if cache is not None else None
             extracted = (deepcopy(previous[1]) if previous and previous[0] == fingerprint
-                         else ADAPTERS[language].extract(batch, max_lines, settings))
+                         else extract_with_fallback(ADAPTERS[language], batch, max_lines, settings))
             if cache is not None:
                 pending_cache[key] = (fingerprint, deepcopy(extracted))
             offset = len(units)
@@ -124,6 +150,9 @@ def source_units(root, files, max_lines=65, language_options=None, report=None,
         cache.clear()
         cache.update(pending_cache)
     if report is not None:
+        diagnostics = [unit['parseDiagnostic'] for unit in units if 'parseDiagnostic' in unit]
         report.update(inputFiles=len(files), acceptedFiles=len(sources), excluded=excluded,
-                      fallbackFiles=sum(language_for(source.path) == 'text' for source in sources))
+                      fallbackFiles=len({source.path for source in sources if language_for(source.path) == 'text'}
+                                        | {item['path'] for item in diagnostics}),
+                      degradedFiles=len(diagnostics), parseDiagnostics=diagnostics)
     return units

package/src/retrieval/languages/go.py

@@ -11,7 +11,7 @@ import subprocess
 import sys
 import tempfile
 
-from .schema import physical_lines
+from .schema import physical_lines, SourceSyntaxError
 
 
 @lru_cache(maxsize=8)
@@ -81,6 +81,8 @@ def parse(sources, options):
 def extract(sources, max_lines=65, options=None):
     options = settings(options)
     parsed = parse(sources, options)
+    if parsed.get('syntaxErrors'):
+        raise SourceSyntaxError(parsed['syntaxErrors'])
     units, records, targets = [], {}, defaultdict(list)
     source_by_path = {source.path: source for source in sources}
     for file in parsed['files']:

package/src/retrieval/languages/go_ast.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"go/ast"
 	"go/parser"
+	"go/scanner"
 	"go/token"
 	"os"
 	"path"
@@ -74,10 +75,17 @@ func run() error {
 	objects := map[*ast.Object]Target{}
 	functions := map[string][]Target{}
 	line := func(p token.Pos) int { return fset.PositionFor(p, false).Line }
+	diagnostics := []map[string]any{}
 	for _, src := range input.Files {
 		text := strings.ReplaceAll(strings.ReplaceAll(src.Text, "\r\n", "\n"), "\r", "\n")
 		tree, err := parser.ParseFile(fset, src.Path, text, parser.ParseComments|parser.AllErrors)
 		if err != nil {
+			if syntax, ok := err.(scanner.ErrorList); ok && len(syntax) > 0 {
+				pos := syntax[0].Pos
+				diagnostics = append(diagnostics, map[string]any{"path": src.Path, "language": "go",
+					"errorType": "SyntaxError", "line": pos.Line, "column": pos.Column})
+				continue
+			}
 			return err
 		}
 		trees[src.Path] = tree
@@ -93,6 +101,9 @@ func run() error {
 			}
 		}
 	}
+	if len(diagnostics) > 0 {
+		return json.NewEncoder(os.Stdout).Encode(map[string]any{"compilerVersion": runtime.Version(), "syntaxErrors": diagnostics})
+	}
 	var typed TypeEvidence
 	if input.Options.Mode == "types" {
 		typed = resolveTypes(fset, trees, texts, input.Options.ModulePath, input.Options.GOOS, input.Options.GOARCH)

package/src/retrieval/languages/python.py

@@ -4,20 +4,32 @@ Name-resolved calls/inheritance remain heuristic (shadowing/dynamic dispatch are
 not fully resolved); confidence is a provenance category, not a probability.
 """
 import ast
+import re
 from collections import defaultdict
+from .schema import SourceSyntaxError
+from .python_calls import resolved_links, symbol_resolver
 
 def extract(sources, max_lines=65, options=None):
     if options:
         raise ValueError("Python adapter does not accept language options")
     units, trees, aliases, class_bases = [], {}, {}, {}
+    diagnostics = []
+    scoped_calls = {}
     for source in sources:
         path, text = source.path, source.text
         lines = text.splitlines()
         module = path.removesuffix('.py').replace('/', '.')
         if module.endswith('.__init__'):
             module = module[:-9]
-        tree = ast.parse(text, filename=path)
+        try:
+            tree = ast.parse(text, filename=path)
+        except SyntaxError as error:
+            diagnostics.append({'path': path, 'language': 'python', 'errorType': type(error).__name__,
+                                'line': error.lineno or 1, 'column': error.offset or 1})
+            continue
         trees[module] = tree
+        package = module if path.endswith('/__init__.py') or path == '__init__.py' else module.rpartition('.')[0]
+        scoped_calls[module] = resolved_links(tree, module, package)
         imports = {}
         for n in tree.body:
             if isinstance(n, ast.Import):
@@ -74,9 +86,12 @@ def extract(sources, max_lines=65, options=None):
 
         definitions(tree.body)
 
+    if diagnostics:
+        raise SourceSyntaxError(diagnostics)
     symbols = defaultdict(list)
     for u in units:
         symbols[u['symbol']].append(u['id'])
+    resolve = symbol_resolver(set(trees), symbols)
     for u in units:
         targets, relations = [], []
 
@@ -86,18 +101,13 @@ def extract(sources, max_lines=65, options=None):
                               'resolution': resolution} for target in ids if target != u['id'])
         if u['owner']:
             relate(symbols.get(u['owner'], []), 'member_of', 1.0, 'syntax')
-        for call in u['calls']:
-            pieces = call.split('.')
-            first = pieces[0]
-            target = None
-            if first in ('self', 'cls') and u['owner']:
-                target = u['owner'] + '.' + '.'.join(pieces[1:])
-            elif first in aliases[u['module']]:
-                target = aliases[u['module']][first] + ('.' + '.'.join(pieces[1:]) if len(pieces) > 1 else '')
-            elif len(pieces) == 1:
-                target = u['module'] + '.' + call
-            if target:
-                relate(symbols.get(target, []), 'calls', .7, 'static-name')
+        for kind, links in scoped_calls[u['module']].items():
+            for line, target in links:
+                if u['start'] <= line <= u['end']:
+                    relate(resolve(target), kind, .7, 'scoped-import')
+        for reference in re.findall(r':(?:func|meth):`~?([A-Za-z_][\w.]*)`', u['text']):
+            targets_ = resolve(reference) or resolve(u['module'] + '.' + reference)
+            relate(targets_, 'references_value', .6, 'explicit-doc-reference')
         for base in class_bases.get(u['owner'] or u['symbol'], []):
             pieces = base.split('.')
             target = aliases[u['module']].get(pieces[0], u['module'] + '.' + pieces[0])