fix: concat logic and improve codecov

Author: gusye1234

prompt_string/string.py

@@ -73,8 +73,12 @@ def __add__(self, other):
     def __truediv__(self, other):
         from .string_chain import PromptChain
 
-        assert isinstance(other, PromptString)
-        return PromptChain([self, other])
+        if isinstance(other, PromptString):
+            return PromptChain([self, other])
+        elif isinstance(other, PromptChain):
+            return PromptChain([self] + other.prompts)
+        else:
+            raise ValueError(f"Invalid type for Prompt Division: {type(other)}")
 
     @to_prompt_string
     def replace(self, old, new, count=-1):

prompt_string/string_chain.py

@@ -41,14 +41,18 @@ def infer_roles(self):
     def roles(self):
         return [p.role for p in self.__prompts]
 
+    @property
+    def prompts(self):
+        return self.__prompts
+
     def __truediv__(self, other):
         if isinstance(other, PromptChain):
             return PromptChain(
-                self.__prompts + other.__prompts, default_start_role=self.__start_role
+                self.prompts + other.prompts, default_start_role=self.__start_role
             )
         elif isinstance(other, PromptString):
             return PromptChain(
-                self.__prompts + [other], default_start_role=self.__start_role
+                self.prompts + [other], default_start_role=self.__start_role
             )
         else:
             raise ValueError(f"Invalid type for PromptChain Division: {type(other)}")

prompt_string/token.py

@@ -12,10 +12,6 @@ def get_decoded_tokens(tokens: list[int]) -> str:
     return USE_ENCODER.decode(tokens)
 
 
-def truncate_string(content: str, max_tokens: int):
-    return get_decoded_tokens(get_encoded_tokens(content)[:max_tokens])
-
-
 def setup_encoder(model: str = "gpt-4o"):
     global USE_ENCODER
     USE_ENCODER = encoding_for_model(model)

readme.md

@@ -6,10 +6,17 @@
     <a href="https://pypi.org/project/prompt-string/">
       <img src="https://img.shields.io/pypi/v/prompt-string.svg">
     </a>
+    <a href="https://github.com/memodb-io/prompt-string/actions/workflows/test.yml" > 
+     <img src="https://github.com/memodb-io/prompt-string/actions/workflows/test.yml/badge.svg"/> 
+		</a>
+    <a href="https://codecov.io/github/memodb-io/prompt-string" > 
+     <img src="https://codecov.io/github/memodb-io/prompt-string/graph/badge.svg?token=kgeW8G0HYW"/> 
+		</a>
 </div>
 
 
 
+
 Prompt is essentially a string, but it should behave somewhat differently from a standard string:
 
 📏 **Length & Slicing**: A prompt string should consider the length in terms of tokens, not characters, and slicing should be done accordingly.
@@ -46,9 +53,9 @@ from prompt_string import P
 
 prompt = P("you're a helpful assistant.")
 
-print("Total token size", len(prompt))
-print("Decoded result of the second token", prompt[2])
-print("The decoded result of first five tokens", prompt[:5])
+print("Total token size:", len(prompt))
+print("Decoded result of the second token:", prompt[2])
+print("The decoded result of first three tokens:", prompt[:3])
 ```
 
 `P` supports some `str` native methods to still return a `P` object:
@@ -60,7 +67,7 @@ print("The decoded result of first five tokens", prompt[:5])
 prompt = P("you're a helpful assistant. {temp}")
 
 print(len(prompt.format(temp="End of instructions")))
-print(len(prompt.replace("{temp}", ""))
+print(len(prompt.replace("{temp}", "")))
 ```
 
 > 🧐 Raise an issue if you think other methods should be supported
@@ -75,7 +82,7 @@ from prompt_string import P
 sp = P("you're a helpful assistant.", role="system")
 up = P("How are you?", role="user")
 
-print(sp.role, up.role, (sp+up).roles)
+print(sp.role, up.role, (sp+up).role)
 print(sp + up)
 
 print(sp.message())
@@ -84,7 +91,7 @@ print(sp.message())
 - role can be `None`, `str` for `P`
 - For single prompt, like `sp`, the role is `str`(*e.g.* `system`) or `None`
 - `sp+up` will concatenate two prompt string and generate a new `P`, whose role will be updated if the latter one has one.
-  - For example, `sp+up`'s role is `user`, `sp+P('Hi')`'s role is `system`
+  - For example, `sp+up`'s role is `user`; `sp+P('Hi')`'s role is `system`
 
 
 - `.message(...)` return a JSON object of this prompt.

tests/test_basic.py

@@ -1,3 +1,4 @@
+import pytest
 from prompt_string import P
 from prompt_string.string_chain import PromptChain
 
@@ -11,6 +12,15 @@ def test_basic_role():
     assert sp.role == "user"
 
 
+def test_basic_index():
+    p = P("you're a helpful assistant.")
+
+    assert str(p[0]) == "you're"
+    assert str(p[1:]) == " a helpful assistant."
+    with pytest.raises(ValueError):
+        p[1.5]
+
+
 def test_basic_print():
     p = P("you're a helpful assistant.", role="user")
     print(p)
@@ -30,8 +40,11 @@ def test_basic_print():
         "content": "you're a helpful assistant.you're a helpful assistant.",
     }
 
-    p += p
-    print(p, p.role)
+    p4 = p + "Good"
+    assert p4.message() == {
+        "role": "user",
+        "content": "you're a helpful assistant.Good",
+    }
 
 
 def test_basic_chain():
@@ -42,6 +55,21 @@ def test_basic_chain():
     assert len(pc) == 3
     assert pc.infer_roles == ["user", "assistant", "user"]
     assert pc[2] is p
+    assert len(pc[::2]) == 2
+
+    empty_pc = pc[100:]
+    assert len(empty_pc) == 0
+    assert empty_pc.infer_roles == []
+
+    pc_new = PromptChain([p, p2], default_start_role="system")
+    assert pc_new.roles == [None, None]
+    assert pc_new.infer_roles == ["system", "user"]
+
+    ppc = pc / pc
+    assert len(ppc) == 6
+
+    pppc = p / pc
+    assert len(pppc) == 4
 
 
 def test_basic_str_methods():