06 Protobuf 与 gRPC 通信

message MemoryItem {
  string id = 1;
  string content = 2;
  string memory_type = 3;
  repeated float embedding = 4;
  string created_at = 5;
  string last_accessed = 6;
  int32 access_count = 7;
  string valid_from = 8;
  string valid_until = 9;
  double trust_score = 10;
  double importance = 11;
  string layer = 12;
  double decay_rate = 13;
  string source_id = 14;
  string causal_parent_id = 15;
  string supersedes_id = 16;
  repeated string entity_refs = 17;
  repeated string tags = 18;
  string deleted_at = 19;
}

message SearchResult {
  MemoryItem item = 1;
  double score = 2;
  repeated string matched_by = 3;
}

service StorageService {
  rpc AddMemory(AddMemoryRequest) returns (AddMemoryResponse);
  rpc GetMemory(GetMemoryRequest) returns (GetMemoryResponse);
  rpc UpdateMemory(UpdateMemoryRequest) returns (UpdateMemoryResponse);
  rpc DeleteMemory(DeleteMemoryRequest) returns (DeletedResponse);
  rpc SearchQuery(SearchQueryRequest) returns (SearchResultList);
  rpc SearchFullText(SearchFullTextRequest) returns (SearchResultList);
  rpc SearchByEntities(SearchByEntitiesRequest) returns (SearchResultList);
  rpc SearchByVector(SearchByVectorRequest) returns (SearchResultList);
  rpc TouchMemory(TouchMemoryRequest) returns (BoolResponse);
  rpc TraceAncestors(TraceAncestorsRequest) returns (MemoryList);
  rpc TraceDescendants(TraceDescendantsRequest) returns (MemoryList);
  rpc ListMemories(ListMemoriesRequest) returns (MemoryList);
  rpc AddRelation(AddRelationRequest) returns (CreatedResponse);
  rpc ListRelations(ListRelationsRequest) returns (RelationList);
  rpc RelationExists(RelationExistsRequest) returns (BoolResponse);
  rpc GetEvolutionEvents(GetEvolutionEventsRequest) returns (EvolutionEventList);
  rpc GetAuditEvents(GetAuditEventsRequest) returns (AuditEventList);
  rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse);
}

proto:
	export PATH="$$(go env GOPATH)/bin:$$PATH" && protoc -I $(PROTO_DIR) \
		--go_out=go-server/gen --go_opt=paths=source_relative \
		--go-grpc_out=go-server/gen --go-grpc_opt=paths=source_relative \
		$(PROTO_DIR)/memory/v1/models.proto \
		$(PROTO_DIR)/memory/v1/storage_service.proto \
		$(PROTO_DIR)/memory/v1/ai_service.proto
	. .venv/bin/activate && python -m grpc_tools.protoc -I $(PROTO_DIR) \
		--python_out=src/agent_memory/generated \
		--grpc_python_out=src/agent_memory/generated \
		$(PROTO_DIR)/memory/v1/models.proto \
		$(PROTO_DIR)/memory/v1/storage_service.proto \
		$(PROTO_DIR)/memory/v1/ai_service.proto

make proto

grpcServer := grpc.NewServer(grpc.UnaryInterceptor(grpcserver.UnaryAuthInterceptor(cfg)))
memoryv1.RegisterStorageServiceServer(grpcServer, grpcserver.New(backend, orchestrator, metrics))

import memoryv1 "github.com/bakebakebakebake/agent-memory/go-server/gen/memory/v1"

func (server *Server) SearchQuery(ctx context.Context, request *memoryv1.SearchQueryRequest) (*memoryv1.SearchResultList, error) {
	start := time.Now()
	results, err := server.orchestrator.Search(ctx, request.Query, request.Embedding, request.Entities, request.Limit)
	if err != nil {
		return nil, err
	}
	server.observeSearch("grpc", "fused", start)
	return &memoryv1.SearchResultList{Results: results}, nil
}

if self.config.prefer_grpc and grpc is not None and storage_service_pb2_grpc is not None:
    self._grpc_channel = grpc.insecure_channel(self.config.grpc_target)
    self._grpc_stub = storage_service_pb2_grpc.StorageServiceStub(self._grpc_channel)

def _memory_to_proto(self, item: MemoryItem):
    if models_pb2 is None:
        raise RuntimeError("gRPC stubs are unavailable. Install remote dependencies and regenerate protos.")
    return models_pb2.MemoryItem(
        id=item.id,
        content=item.content,
        memory_type=item.memory_type.value,
        embedding=item.embedding,
        created_at=item.created_at.isoformat(),
        last_accessed=item.last_accessed.isoformat(),
        access_count=item.access_count,
        valid_from=item.valid_from.isoformat() if item.valid_from else "",
        valid_until=item.valid_until.isoformat() if item.valid_until else "",
        trust_score=item.trust_score,
        importance=item.importance,
        layer=item.layer.value,
        decay_rate=item.decay_rate,
        source_id=item.source_id,
        causal_parent_id=item.causal_parent_id or "",
        supersedes_id=item.supersedes_id or "",
        entity_refs=item.entity_refs,
        tags=item.tags,
        deleted_at=item.deleted_at.isoformat() if item.deleted_at else "",
    )

def _memory_from_proto(self, item) -> MemoryItem:
    return _memory_from_payload(
        {
            "id": item.id,
            "content": item.content,
            "memory_type": item.memory_type,
            "embedding": list(item.embedding),
            "created_at": item.created_at,
            "last_accessed": item.last_accessed,
            "access_count": item.access_count,
            "valid_from": item.valid_from or None,
            "valid_until": item.valid_until or None,
            "trust_score": item.trust_score,
            "importance": item.importance,
            "layer": item.layer,
            "decay_rate": item.decay_rate,
            "source_id": item.source_id,
            "causal_parent_id": item.causal_parent_id or None,
            "supersedes_id": item.supersedes_id or None,
            "entity_refs": list(item.entity_refs),
            "tags": list(item.tags),
            "deleted_at": item.deleted_at or None,
        }
    )

def add_memory(self, item: MemoryItem) -> MemoryItem:
    if self._grpc_stub is not None:
        response = self._grpc_call("AddMemory", storage_service_pb2.AddMemoryRequest(item=self._memory_to_proto(item)))
        return self._memory_from_proto(response.item)
    payload = self._request_json("POST", "/api/v1/memories", data=_memory_to_payload(item))
    return _memory_from_payload(payload["item"])

sequenceDiagram
    participant C as MemoryClient
    participant R as RemoteBackend
    participant S as gRPC Stub
    participant I as UnaryAuthInterceptor
    participant G as Go StorageService

    C->>R: add(item)
    R->>R: _memory_to_proto(item)
    R->>S: AddMemory(request, metadata)
    S->>I: unary call
    I->>G: auth ok 后继续
    G->>G: backend.AddMemory(...)
    G-->>S: AddMemoryResponse
    S-->>R: response
    R->>R: _memory_from_proto(response.item)
    R-->>C: MemoryItem

def _grpc_call(self, method_name: str, request_message):
    if self._grpc_stub is None:
        raise RuntimeError("gRPC stub is unavailable")
    metadata: list[tuple[str, str]] = []
    if self.config.api_key:
        metadata.append(("x-api-key", self.config.api_key))
    if self.config.jwt_token:
        metadata.append(("authorization", f"Bearer {self.config.jwt_token}"))
    method = getattr(self._grpc_stub, method_name)
    return method(request_message, metadata=metadata or None)

func UnaryAuthInterceptor(cfg config.Config) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, request any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
		if cfg.APIKey == "" && cfg.JWTSecret == "" {
			return handler(ctx, request)
		}
		md, _ := metadata.FromIncomingContext(ctx)
		apiKey := firstValue(md.Get("x-api-key"))
		authorization := firstValue(md.Get("authorization"))
		if cfg.APIKey != "" && auth.APIKeyMatches(apiKey, cfg.APIKey) {
			return handler(ctx, request)
		}
		if cfg.JWTSecret != "" && auth.JWTMatches(auth.ParseBearerToken(authorization), cfg.JWTSecret) {
			return handler(ctx, request)
		}
		return nil, status.Error(codes.Unauthenticated, "unauthorized")
	}
}